Merge branch 'main' of github.com:Ostorlab/ostorlab into feature/add-…

…support-for-stack-traces-in-vulnerability-location-proto-msg

setup.cfg

@@ -13,7 +13,7 @@ license_files = LICENSE
 long_description = file: README.md
 long_description_content_type = text/markdown
 url = https://github.com/Ostorlab/ostorlab
-version=1.4.3
+version=1.4.4
 project_urls =
     Documentation = https://oxo.ostorlab.co/
     Source = https://github.com/Ostorlab/oxo

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 @@ def _get(key):
         if key in storage:
             return storage[key]
 
-    def _add(key, value):
+    def _add(key: Union[str, bytes], value: bytes) -> None:
         """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 @@ def _value_type(key):
         "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