Update agent_persist_mock to patch _exists method

Ostorlab · Feb 23, 2025 · c710b37 · c710b37
1 parent 776043e
commit c710b37
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/ostorlab/testing/agent.py b/src/ostorlab/testing/agent.py
@@ -97,7 +97,7 @@ def _get(key):
     def _add(key: Union[str, bytes], value: bytes):
         """Check values are present in the storage dict."""
         if isinstance(value, bytes) is False:
-            storage[key] = str(value).encode()
+            value = str(value).encode()
         storage[key] = value
 
     def _exists(key: str) -> bool:

diff --git a/tests/testing/agent_test.py b/tests/testing/agent_test.py
@@ -62,3 +62,6 @@ def testMockPersistAgent_whensetMethodsAreCalled_stateIsPersistedByMock(
     assert test_agent.set_add("test", "1") is True
     assert test_agent.set_is_member("test", "1") is True
     assert agent_persist_mock == {"test": {"1"}}
+    assert test_agent.exists("akey") is False
+    test_agent.add("akey", b"aval")
+    assert test_agent.exists("akey") is True