1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-11-22 15:38:11 +00:00
trezor-firmware/core/tests/mock_storage.py
obrusvit e073e619c9 chore(tests): re-run black and isort on core/tests
isort set to skip the first necessary "from common import *" line. A
better solution would be to get rid of the need of this import in the
future.

[no changelog]
2024-02-22 12:10:12 +01:00

43 lines
1.2 KiB
Python

import storage.common
from mock import patch
class MockStorage:
PATCH_METHODS = ("get", "set", "delete")
def __init__(self):
self.namespace = {}
self.patches = [
patch(storage.common, method, getattr(self, method))
for method in self.PATCH_METHODS
]
def set(self, app: int, key: int, data: bytes, public: bool = False) -> None:
self.namespace.setdefault(app, {})
self.namespace[app][key] = data
def get(self, app: int, key: int, public: bool = False) -> bytes | None:
self.namespace.setdefault(app, {})
return self.namespace[app].get(key)
def delete(self, app: int, key: int, public: bool = False) -> None:
self.namespace.setdefault(app, {})
self.namespace[app].pop(key, None)
def __enter__(self):
for self_patch in self.patches:
self_patch.__enter__()
return self
def __exit__(self, exc_type, exc_value, tb):
for self_patch in self.patches:
self_patch.__exit__(exc_type, exc_value, tb)
def mock_storage(func):
def inner(*args, **kwargs):
with MockStorage():
return func(*args, **kwargs)
return inner