minor test fix

Signed-off-by: Lorenzo Curcio <Lorenzo.curcio@valuecuberesearch.com>

Author: Lorenzo Curcio

dapr/actor/runtime/mock_state_manager.py

@@ -183,8 +183,10 @@ async def add_or_update_state(
             self._default_state_change_tracker[state_name] = StateMetadata(
                 new_value, StateChangeKind.update
             )
+            self._mock_state[state_name] = new_value
             return new_value
         self._default_state_change_tracker[state_name] = StateMetadata(value, StateChangeKind.add)
+        self._mock_state[state_name] = value
         return value
 
     async def get_state_names(self) -> List[str]:

tests/actor/test_mock_actor.py

@@ -24,6 +24,14 @@ async def clear_data(self) -> None:
     async def test_data(self) -> int:
         ...
 
+    @actormethod(name='AddState')
+    async def add_state(self, name: str, data: object) -> None:
+        ...
+
+    @actormethod(name='UpdateState')
+    async def update_state(self, name: str, data: object) -> None:
+        ...
+
     @actormethod(name='AddDataNoSave')
     async def add_data_no_save(self, data: object) -> None:
         ...
@@ -77,6 +85,15 @@ async def test_data(self) -> int:
             return 3
         return 4
 
+    async def add_state(self, name: str, data: object) -> None:
+        await self._state_manager.add_state(name, data)
+
+    async def update_state(self, name: str, data: object) -> None:
+        def double(_: str, x: int) -> int:
+            return 2 * x
+
+        await self._state_manager.add_or_update_state(name, data, double)
+
     async def add_data_no_save(self, data: object) -> None:
         await self._state_manager.set_state('state', data)
 
@@ -227,6 +244,31 @@ async def test_test_data(self):
         result = await mockactor.test_data()
         self.assertEqual(result, 3)
 
+    async def test_add_state(self):
+        mockactor = create_mock_actor(MockTestActor, '1')
+        print(mockactor._state_manager._mock_state)
+        self.assertFalse(mockactor._state_manager._mock_state)
+        await mockactor.add_state('test', 5)
+        self.assertTrue('test' in mockactor._state_manager._mock_state)
+        self.assertEqual(mockactor._state_manager._mock_state['test'], 5)
+        await mockactor.add_state('test2', 10)
+        self.assertTrue('test2' in mockactor._state_manager._mock_state)
+        self.assertEqual(mockactor._state_manager._mock_state['test2'], 10)
+        self.assertEqual(len(mockactor._state_manager._mock_state), 2)
+        with self.assertRaises(ValueError):
+            await mockactor.add_state('test', 10)
+
+    async def test_update_state(self):
+        mockactor = create_mock_actor(MockTestActor, '1')
+        self.assertFalse(mockactor._state_manager._mock_state)
+        await mockactor.update_state('test', 10)
+        self.assertTrue('test' in mockactor._state_manager._mock_state)
+        self.assertEqual(mockactor._state_manager._mock_state['test'], 10)
+        await mockactor.update_state('test', 10)
+        self.assertTrue('test' in mockactor._state_manager._mock_state)
+        self.assertEqual(mockactor._state_manager._mock_state['test'], 20)
+        self.assertEqual(len(mockactor._state_manager._mock_state), 1)
+
     async def test_state_change_tracker(self):
         mockactor = create_mock_actor(MockTestActor, '1')
         self.assertEqual(len(mockactor._state_manager._default_state_change_tracker), 0)