Avoid set several times HACachedStates (#56)

* Avoid set several times HACachedStates * Improve testing
home-assistant · Mar 19, 2024 · 8d1361d · 8d1361d
1 parent 48e36a0
commit 8d1361d
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 26 deletions.
diff --git a/Source/Caches/HACacheKeyStates.swift b/Source/Caches/HACacheKeyStates.swift
@@ -24,25 +24,26 @@ internal struct HACacheKeyStates: HACacheKey {
     ///   - shouldResetEntities: True if current state needs to be ignored (e.g. re-connection)
     /// - Returns: HAEntity cached states
     /// Logic from: https://github.com/home-assistant/home-assistant-js-websocket/blob/master/lib/entities.ts
+    // swiftlint: disable cyclomatic_complexity
     static func processUpdates(
         info: HACacheTransformInfo<HACompressedStatesUpdates, HACachedStates?>,
         shouldResetEntities: Bool
-    )
-        -> HACachedStates {
-        var states: HACachedStates
-        if shouldResetEntities {
-            states = .init(entities: [])
-        } else {
-            states = info.current ?? .init(entities: [])
+    ) -> HACachedStates {
+        var updatedEntities: [String: HAEntity] = [:]
+
+        if !shouldResetEntities, let currentEntities = info.current {
+            updatedEntities = currentEntities.allByEntityId
         }
+
         if let additions = info.incoming.add {
             for (entityId, updates) in additions {
-                if var currentState = states[entityId] {
+                if var currentState = updatedEntities[entityId] {
                     currentState.update(from: updates)
-                    states[entityId] = currentState
+                    updatedEntities[entityId] = currentState
                 } else {
                     do {
-                        states[entityId] = try updates.asEntity(entityId: entityId)
+                        let newEntity = try updates.asEntity(entityId: entityId)
+                        updatedEntities[entityId] = newEntity
                     } catch {
                         HAGlobal.log(.error, "[Update-To-Entity-Error] Failed adding new entity: \(error)")
                     }
@@ -52,26 +53,26 @@ internal struct HACacheKeyStates: HACacheKey {
 
         if let subtractions = info.incoming.remove {
             for entityId in subtractions {
-                states[entityId] = nil
+                updatedEntities.removeValue(forKey: entityId)
             }
         }
 
         if let changes = info.incoming.change {
-            changes.forEach { entityId, diff in
-                guard var entityState = states[entityId] else { return }
+            for (entityId, diff) in changes {
+                guard var entityState = updatedEntities[entityId] else { continue }
 
                 if let toAdd = diff.additions {
                     entityState.add(toAdd)
-                    states[entityId] = entityState
                 }
 
                 if let toRemove = diff.subtractions {
                     entityState.subtract(toRemove)
-                    states[entityId] = entityState
                 }
+
+                updatedEntities[entityId] = entityState
             }
         }
 
-        return states
+        return .init(entitiesDictionary: updatedEntities)
     }
 }
diff --git a/Source/Caches/HACachedStates.swift b/Source/Caches/HACachedStates.swift
@@ -6,23 +6,31 @@ public extension HACachesContainer {
 /// Cached version of all entity states
 public struct HACachedStates {
     /// All entities
-    public var all: Set<HAEntity>
+    public var all: Set<HAEntity> = []
     /// All entities, keyed by their entityId
     public subscript(entityID: String) -> HAEntity? {
         get { allByEntityId[entityID] }
         set { allByEntityId[entityID] = newValue }
     }
 
     /// Backing dictionary, whose mutation updates the set
-    private var allByEntityId: [String: HAEntity] {
+    internal var allByEntityId: [String: HAEntity] {
         didSet {
             all = Set(allByEntityId.values)
         }
     }
 
     /// Create a cached state
+    /// - Parameter entitiesDictionary: The entities to start with, key is the entity ID
+    public init(entitiesDictionary: [String: HAEntity]) {
+        self.all = Set(entitiesDictionary.values)
+        self.allByEntityId = entitiesDictionary
+    }
+
+    /// Create a cached state
+    /// Mainly for tests
     /// - Parameter entities: The entities to start with
-    public init(entities: [HAEntity]) {
+    internal init(entities: [HAEntity]) {
         self.all = Set(entities)
         self.allByEntityId = entities.reduce(into: [:]) { dictionary, entity in
             dictionary[entity.entityId] = entity

diff --git a/Source/Data/HAEntity.swift b/Source/Data/HAEntity.swift
@@ -53,7 +53,7 @@ public struct HAEntity: HADataDecodable, Hashable {
         attributes: [String: Any],
         context: HAResponseEvent.Context
     ) throws {
-        var domain: String = try {
+        let domain: String = try {
             if let domain {
                 domain
             } else {

diff --git a/Tests/HACacheKeyStates.test.swift b/Tests/HACacheKeyStates.test.swift
@@ -55,7 +55,7 @@ internal final class HACacheKeyStates_test: XCTestCase {
                     )
                 ),
                 current: .init(
-                    entities: []
+                    entitiesDictionary: [:]
                 )
             ), shouldResetEntities: false
         )
@@ -133,8 +133,8 @@ internal final class HACacheKeyStates_test: XCTestCase {
                     )
                 ),
                 current: .init(
-                    entities: [
-                        existentEntity,
+                    entitiesDictionary: [
+                        "person.bruno": existentEntity,
                     ]
                 )
             ), shouldResetEntities: false
@@ -202,8 +202,8 @@ internal final class HACacheKeyStates_test: XCTestCase {
                     )
                 ),
                 current: .init(
-                    entities: [
-                        existentEntity,
+                    entitiesDictionary: [
+                        "person.bruno": existentEntity,
                     ]
                 )
             ), shouldResetEntities: false
@@ -251,7 +251,7 @@ internal final class HACacheKeyStates_test: XCTestCase {
                     )
                 ),
                 current: .init(
-                    entities: []
+                    entitiesDictionary: [:]
                 )
             ), shouldResetEntities: false
         )

diff --git a/Tests/HACachedStates.test.swift b/Tests/HACachedStates.test.swift
@@ -125,4 +125,27 @@ internal class HACachedStatesTests: XCTestCase {
 
         XCTAssertEqual(entityRemovedOutgoingType?.all.count, 0)
     }
+
+    func testSubscriptByEntityIdReturnsCorrectly() throws {
+        let expectedEntity: HAEntity = try .fake(id: "person.bruno")
+        let cache = HACachedStates(entitiesDictionary: [expectedEntity.entityId: expectedEntity])
+        XCTAssertEqual(cache["fake.person.bruno"], expectedEntity)
+    }
+
+    func testSubscriptSetByEntityIdSetsCorrectly() throws {
+        let expectedEntity: HAEntity = try .fake(id: "person.bruno")
+        let expectedEntity2: HAEntity = try .fake(id: "person.bruno2")
+        var cache = HACachedStates(entitiesDictionary: [expectedEntity.entityId: expectedEntity])
+        cache[expectedEntity2.entityId] = expectedEntity2
+        XCTAssertEqual(cache["fake.person.bruno2"], expectedEntity2)
+    }
+
+    func testSetAllByEntityIdSetsAllToo() throws {
+        let expectedEntity: HAEntity = try .fake(id: "person.bruno")
+        let expectedEntity2: HAEntity = try .fake(id: "person.bruno2")
+        var cache = HACachedStates(entitiesDictionary: [expectedEntity.entityId: expectedEntity])
+
+        cache.allByEntityId = [expectedEntity2.entityId: expectedEntity2]
+        XCTAssertEqual(cache.allByEntityId, [expectedEntity2.entityId: expectedEntity2])
+    }
 }