Skip to content

Commit

Permalink
Update agent_persist_mock to patch _exists method
Browse files Browse the repository at this point in the history
  • Loading branch information
Mohamed Nasser committed Feb 23, 2025
1 parent 2a0f88e commit 376e4bd
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions src/ostorlab/testing/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,9 +95,15 @@ def _get(key):
if key in storage:
return storage[key]

def _add(key, value):
def _add(key: str | bytes, value: bytes):
"""Check values are present in the storage dict."""
storage[key] = str(value).encode()
if isinstance(value, bytes) is False:
storage[key] = str(value).encode()
storage[key] = value

Check warning on line 102 in src/ostorlab/testing/agent.py

View check run for this annotation

Codecov / codecov/patch

src/ostorlab/testing/agent.py#L100-L102

Added lines #L100 - L102 were not covered by tests

def _exists(key: str) -> bool:
"""Check if thr key is present in the storage dict."""
return key in storage

Check warning on line 106 in src/ostorlab/testing/agent.py

View check run for this annotation

Codecov / codecov/patch

src/ostorlab/testing/agent.py#L106

Added line #L106 was not covered by tests

def _hash_add(hash_name, mapping):
for k, v in mapping.items():
Expand Down Expand Up @@ -156,6 +162,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,
Expand Down

0 comments on commit 376e4bd

Please sign in to comment.