Support get_state namespaces

appdaemontestframework/given_that.py

@@ -1,16 +1,24 @@
+from collections import defaultdict
 from datetime import datetime
 
 from appdaemontestframework.common import AppdaemonTestFrameworkError
 from appdaemontestframework.hass_mocks import HassMocks
 
 
 class StateNotSetError(AppdaemonTestFrameworkError):
-    def __init__(self, entity_id):
-        super().__init__(f"""
-        State for entity: '{entity_id}' was never set!
-        Please make sure to set the state with `given_that.state_of({entity_id}).is_set_to(STATE)`
-        before trying to access the mocked state
-        """)
+    def __init__(self, entity_id, namespace):
+        if namespace != 'default':
+            super().__init__(f"""
+            State for entity: '{entity_id}' in '{namespace}' namespace was never set!
+            Please make sure to set the state with `given_that.state_of({entity_id}, NAMESPACE).is_set_to(STATE)`
+            before trying to access the mocked state
+            """)
+        else:
+            super().__init__(f"""
+            State for entity: '{entity_id}' was never set!
+            Please make sure to set the state with `given_that.state_of({entity_id}).is_set_to(STATE)`
+            before trying to access the mocked state
+            """)
 
 
 class AttributeNotSetError(AppdaemonTestFrameworkError):
@@ -24,13 +32,14 @@ def __init__(self, hass_mocks: HassMocks):
         self._init_mocked_passed_args()
 
     def _init_mocked_states(self):
-        self.mocked_states = {}
+        self.mocked_states = defaultdict(dict)
 
-        def get_state_mock(entity_id=None, *, attribute=None):
+        def get_state_mock(entity_id=None, *, attribute=None, namespace=None):
+            namespace = namespace or "default"
             if entity_id is None:
                 resdict = dict()
-                for entityid in self.mocked_states:
-                    state = self.mocked_states[entityid]
+                for entityid in self.mocked_states[namespace]:
+                    state = self.mocked_states[namespace][entityid]
                     resdict.update({
                         entityid: {
                             "state": state['main'],
@@ -39,10 +48,10 @@ def get_state_mock(entity_id=None, *, attribute=None):
                     })
                 return resdict
             else:
-                if entity_id not in self.mocked_states:
-                    raise StateNotSetError(entity_id)
+                if entity_id not in self.mocked_states[namespace]:
+                    raise StateNotSetError(entity_id, namespace)
 
-                state = self.mocked_states[entity_id]
+                state = self.mocked_states[namespace][entity_id]
 
                 if attribute is None:
                     return state['main']
@@ -63,16 +72,18 @@ def format_time(timestamp: datetime):
 
         self._hass_mocks.hass_functions['get_state'].side_effect = get_state_mock
 
-        def entity_exists_mock(entity_id):
-            return entity_id in self.mocked_states
+        def entity_exists_mock(entity_id, namespace=None):
+            namespace = namespace or "default"
+            return entity_id in self.mocked_states[namespace]
 
         self._hass_mocks.hass_functions['entity_exists'].side_effect = entity_exists_mock
 
     def _init_mocked_passed_args(self):
         self.mocked_passed_args = self._hass_mocks.hass_functions['args']
         self.mocked_passed_args.clear()
 
-    def state_of(self, entity_id):
+    def state_of(self, entity_id, namespace=None):
+        namespace = namespace or "default"
         given_that_wrapper = self
 
         class IsWrapper:
@@ -83,7 +94,7 @@ def is_set_to(self,
                           last_changed: datetime = None):
                 if not attributes:
                     attributes = {}
-                given_that_wrapper.mocked_states[entity_id] = {
+                given_that_wrapper.mocked_states[namespace][entity_id] = {
                     'main': state,
                     'attributes': attributes,
                     'last_updated': last_updated,

test/integration_tests/apps/kitchen.py

@@ -18,6 +18,12 @@
 MSG_LONG_OFF = f"was turned off for {LONG_DELAY} minutes"
 MSG_ON = "was turned back on"
 
+STORAGE_NAMESPACE = "app_storage"
+
+PRESETS = {
+    "BRIGHT": {"brightness": 100},
+    "DARK": {"brightness": 20},
+}
 
 class Kitchen(hass.Hass):
     def initialize(self):
@@ -47,7 +53,12 @@ def _new_button_double_click(self, _e, _d, _k):
         self._send_water_heater_notification(MSG_LONG_OFF)
 
     def _new_button_long_press(self, _e, _d, _k):
-        pass
+        preset = PRESETS.get(
+            self.get_state(
+                ID["kitchen"]["light"], attribute="preset", namespace=STORAGE_NAMESPACE
+            )
+        )
+        self.turn_on(ID['kitchen']['light'], **preset)
 
     def _turn_off_water_heater_for_X_minutes(self, minutes):
         self.turn_off(ID['bathroom']['water_heater'])

test/integration_tests/tests/test_kitchen.py

@@ -1,4 +1,4 @@
-from apps.kitchen import Kitchen
+from apps.kitchen import STORAGE_NAMESPACE, Kitchen
 import pytest
 from mock import patch, MagicMock
 from apps.entity_ids import ID
@@ -258,3 +258,14 @@ def test_long_then_short_keep_latest(self, when_new, time_travel, assert_that):
             assert_that(ID['bathroom']['water_heater']).was_not.turned_on()
             time_travel.fast_forward(1).minutes()
             assert_that(ID['bathroom']['water_heater']).was.turned_on()
+
+
+class TestStateStorage:
+    @pytest.mark.parametrize("preset,brightness", [("DARK", 20), ("BRIGHT", 100)])
+    def test_turn_on_preset_light(self, preset, brightness, given_that, when_new, assert_that):
+        given_that.state_of(ID["kitchen"]["light"], namespace=STORAGE_NAMESPACE).is_set_to(
+            "ok", attributes={"preset": preset}
+        )
+
+        when_new.click_button(type="long")
+        assert_that(ID["kitchen"]["light"]).was.turned_on(brightness=brightness)

test/test_state.py

@@ -15,20 +15,20 @@ class MockAutomation(Hass):
     def initialize(self):
         pass
 
-    def is_light_turned_on(self) -> bool:
-        return self.get_state(LIGHT) == 'on'
+    def is_light_turned_on(self, namespace=None) -> bool:
+        return self.get_state(LIGHT, namespace=namespace) == 'on'
 
-    def get_light_brightness(self) -> int:
-        return self.get_state(LIGHT, attribute='brightness')
+    def get_light_brightness(self, namespace=None) -> int:
+        return self.get_state(LIGHT, attribute='brightness', namespace=namespace)
 
-    def get_all_attributes_from_light(self):
-        return self.get_state(LIGHT, attribute='all')
+    def get_all_attributes_from_light(self, namespace=None):
+        return self.get_state(LIGHT, attribute='all', namespace=namespace)
 
     def get_without_using_keyword(self):
         return self.get_state(LIGHT, 'brightness')
 
-    def get_complete_state_dictionary(self):
-        return self.get_state()
+    def get_complete_state_dictionary(self, namespace=None):
+        return self.get_state(namespace=namespace)
 
 
 @automation_fixture(MockAutomation)
@@ -41,6 +41,9 @@ def test_state_was_never_set__raise_error(given_that,
     with raises(StateNotSetError, match=r'.*State.*was never set.*'):
         automation.get_light_brightness()
 
+    with raises(StateNotSetError, match=r'.*State.*namespace was never set.*'):
+        automation.get_light_brightness(namespace='test')
+
 
 def test_set_and_get_state(given_that, automation: MockAutomation):
     given_that.state_of(LIGHT).is_set_to('off')
@@ -49,12 +52,21 @@ def test_set_and_get_state(given_that, automation: MockAutomation):
     given_that.state_of(LIGHT).is_set_to('on')
     assert automation.is_light_turned_on()
 
+    given_that.state_of(LIGHT, namespace='test').is_set_to('off')
+    assert not automation.is_light_turned_on(namespace='test')
+
+    given_that.state_of(LIGHT, namespace='test').is_set_to('on')
+    assert automation.is_light_turned_on(namespace='test')
+
 
 def test_attribute_was_never_set__raise_error(given_that,
                                               automation: MockAutomation):
     given_that.state_of(LIGHT).is_set_to('on')
     assert automation.get_light_brightness() is None
 
+    given_that.state_of(LIGHT, namespace='test').is_set_to('on')
+    assert automation.get_light_brightness(namespace='test') is None
+
 
 def test_set_and_get_attribute(given_that, automation: MockAutomation):
     given_that.state_of(LIGHT).is_set_to('on', attributes={'brightness': 11})
@@ -63,14 +75,33 @@ def test_set_and_get_attribute(given_that, automation: MockAutomation):
     given_that.state_of(LIGHT).is_set_to('on', {'brightness': 22})
     assert automation.get_light_brightness() == 22
 
+    given_that.state_of(LIGHT, namespace='test').is_set_to('on', attributes={'brightness': 11})
+    assert automation.get_light_brightness(namespace='test') == 11
+
+    given_that.state_of(LIGHT, namespace='test').is_set_to('on', {'brightness': 22})
+    assert automation.get_light_brightness(namespace='test') == 22
+
 
 def test_set_and_get_all_attribute(given_that, automation: MockAutomation):
-    given_that.state_of(LIGHT).is_set_to('on', attributes={'brightness': 11,
-                                                           'color': 'blue'})
+    given_that.state_of(LIGHT).is_set_to(
+        "on", attributes={"brightness": 11, "color": "blue"}
+    )
+    given_that.state_of(LIGHT, namespace="test").is_set_to(
+        "on", attributes={"brightness": 22, "color": "red"}
+    )
     assert automation.get_all_attributes_from_light() == {
-        'state': 'on', 'last_updated': None, 'last_changed': None,
-        'entity_id': LIGHT,
-        'attributes': {'brightness': 11, 'color': 'blue'}
+        "state": "on",
+        "last_updated": None,
+        "last_changed": None,
+        "entity_id": LIGHT,
+        "attributes": {"brightness": 11, "color": "blue"},
+    }
+    assert automation.get_all_attributes_from_light(namespace="test") == {
+        "state": "on",
+        "last_updated": None,
+        "last_changed": None,
+        "entity_id": LIGHT,
+        "attributes": {"brightness": 22, "color": "red"},
     }