Update agent_persist_mock to patch _exists method

src/ostorlab/testing/agent.py

@@ -1,11 +1,10 @@
 """mock agent implements the required methods to test the agent's behavior without using external components."""
 
+from typing import List, Union
 import dataclasses
 
 import pytest
 
-from typing import List
-
 from ostorlab.agent.message import message as msg
 
 
@@ -95,9 +94,15 @@
         if key in storage:
             return storage[key]
 
-    def _add(key, value):
+    def _add(key: Union[str, bytes], value: bytes):
         """Check values are present in the storage dict."""
-        storage[key] = str(value).encode()
+        if isinstance(value, bytes) is False:
+            value = str(value).encode()
+        storage[key] = value
+
+    def _exists(key: Union[str, bytes]) -> bool:
+        """Check if thr key is present in the storage dict."""
+        return key in storage
 
     def _hash_add(hash_name, mapping):
         for k, v in mapping.items():
@@ -156,6 +161,10 @@
         "ostorlab.agent.mixins.agent_persist_mixin.AgentPersistMixin.add",
         side_effect=_add,
     )
+    mocker.patch(
+        "ostorlab.agent.mixins.agent_persist_mixin.AgentPersistMixin.exists",
+        side_effect=_exists,
+    )
     mocker.patch(
         "ostorlab.agent.mixins.agent_persist_mixin.AgentPersistMixin.hash_add",
         side_effect=_hash_add,

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