Merge pull request #2491 from fjarri/in-memory-datastore

In-memory mock LMDB for tests

tests/conftest.py

@@ -18,6 +18,7 @@ along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
 
 from collections import defaultdict
 
+import lmdb
 import pytest
 import tempfile
 from pathlib import Path
@@ -29,6 +30,7 @@ from nucypher.network.trackers import AvailabilityTracker
 from nucypher.policy.identity import Card
 from nucypher.utilities.logging import GlobalLoggerSettings
 from tests.constants import INSECURE_DEVELOPMENT_PASSWORD
+from tests.mock.datastore import mock_lmdb_open
 
 # Crash on server error by default
 WebEmitter._crash_on_error_default = True
@@ -77,8 +79,8 @@ def __very_pretty_and_insecure_scrypt_do_not_use():
     Scrypt.derive = original_derivation_function
 
 
-@pytest.fixture(scope='module')
-def monkeymodule():
+@pytest.fixture(scope='session')
+def monkeysession():
     from _pytest.monkeypatch import MonkeyPatch
     mpatch = MonkeyPatch()
     yield mpatch
@@ -199,3 +201,9 @@ def patch_card_directory(session_mocker):
                                 new_callable=session_mocker.PropertyMock)
     yield
     tmpdir.cleanup()
+
+
+@pytest.fixture(scope='session', autouse=True)
+def mock_datastore(monkeysession):
+    monkeysession.setattr(lmdb, 'open', mock_lmdb_open)
+    yield

tests/constants.py

@@ -119,8 +119,6 @@ MOCK_POLICY_DEFAULT_M = 3
 MOCK_IP_ADDRESS = '192.0.2.100'
 MOCK_IP_ADDRESS_2 = '203.0.113.20'
 
-MOCK_URSULA_DB_FILEPATH = tempfile.mkdtemp()
-
 FEE_RATE_RANGE = (5, 10, 15)
 
 

tests/integration/config/test_storages.py

@@ -17,15 +17,16 @@ along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
 
 import os
 import pytest
+import tempfile
 
 from nucypher.characters.lawful import Ursula
 from nucypher.config.storages import ForgetfulNodeStorage, NodeStorage, TemporaryFileBasedNodeStorage
 from nucypher.network.nodes import Learner
 
-from tests.constants import MOCK_URSULA_DB_FILEPATH
 from tests.utils.ursula import MOCK_URSULA_STARTING_PORT
 
 ADDITIONAL_NODES_TO_LEARN_ABOUT = 10
+MOCK_URSULA_DB_FILEPATH = tempfile.mkdtemp()
 
 
 class BaseTestNodeStorageBackends:

tests/integration/conftest.py

@@ -18,10 +18,8 @@ along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
 import lmdb
 import os
 import pytest
-import threading
 from contextlib import contextmanager
 from eth_account.account import Account
-from functools import partial
 
 from nucypher.blockchain.economics import EconomicsFactory
 from nucypher.blockchain.eth.agents import (
@@ -55,42 +53,6 @@ from tests.utils.config import (
 from tests.utils.ursula import MOCK_URSULA_STARTING_PORT
 
 
-class TestLMDBEnv:
-    """
-    This class is used to have LMDB environments open just-in-time as they're
-    needed.
-
-    This is necessary in testing environments because there may be too many
-    LMDB environments open at once.
-    """
-    __test__ = False    # Prohibit pytest from collecting this
-
-    LMDB_OPEN_FUNC = lmdb.open
-    THREAD_LOCK = threading.Lock()
-
-    def __init__(self, *args, **kwargs):
-        self.db_path = args[0]
-        self.open = partial(self.LMDB_OPEN_FUNC, *args, **kwargs)
-
-    @contextmanager
-    def begin(self, *args, **kwargs):
-        try:
-            with self.THREAD_LOCK:
-                with self.open() as lmdb_env:
-                    with lmdb_env.begin(*args, **kwargs) as lmdb_tx:
-                        yield lmdb_tx
-        finally:
-            pass
-
-    def path(self):
-        return self.db_path
-
-
-@pytest.fixture(autouse=True)
-def JIT_lmdb_env(monkeypatch):
-    monkeypatch.setattr("lmdb.open", TestLMDBEnv)
-
-
 @pytest.fixture(scope='function', autouse=True)
 def mock_contract_agency(monkeypatch, module_mocker, token_economics):
     monkeypatch.setattr(ContractAgency, 'get_agent', MockContractAgency.get_agent)

tests/mock/datastore.py

@@ -0,0 +1,128 @@
+"""
+ This file is part of nucypher.
+
+ nucypher is free software: you can redistribute it and/or modify
+ it under the terms of the GNU Affero General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ nucypher is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ GNU Affero General Public License for more details.
+
+ You should have received a copy of the GNU Affero General Public License
+ along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+from bisect import bisect_left
+from contextlib import contextmanager
+import lmdb
+from threading import Lock
+
+from constant_sorrow.constants import MOCK_DB
+
+
+def mock_lmdb_open(db_path, map_size=10485760):
+    if db_path == MOCK_DB:
+        return MockEnvironment()
+    else:
+        return lmdb.Environment(db_path, map_size=map_size)
+
+
+class MockEnvironment:
+
+    def __init__(self):
+        self._storage = {}
+        self._lock = Lock()
+
+    @contextmanager
+    def begin(self, write=False):
+        with self._lock:
+            with MockTransaction(self, write=write) as tx:
+                yield tx
+
+
+class MockTransaction:
+
+    def __init__(self, env, write=False):
+        self._env = env
+        self._storage = dict(env._storage)
+        self._write = write
+        self._invalid = False
+
+    def __enter__(self):
+        if self._invalid:
+            raise lmdb.Error()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type:
+            self.abort()
+        else:
+            self.commit()
+
+    def put(self, key, value, overwrite=True):
+        if self._invalid:
+            raise lmdb.Error()
+        assert self._write
+        if not overwrite and key in self._storage:
+            return False
+        self._storage[key] = value
+        return True
+
+    def get(self, key, default=None):
+        if self._invalid:
+            raise lmdb.Error()
+        return self._storage.get(key, default)
+
+    def delete(self, key):
+        if self._invalid:
+            raise lmdb.Error()
+        assert self._write
+        if key in self._storage:
+            del self._storage[key]
+            return True
+        else:
+            return False
+
+    def commit(self):
+        if self._invalid:
+            raise lmdb.Error()
+        self._invalidate()
+        self._env._storage = self._storage
+
+    def abort(self):
+        self._invalidate()
+        self._storage = self._env._storage
+
+    def _invalidate(self):
+        self._invalid = True
+
+    def cursor(self):
+        return MockCursor(self)
+
+
+class MockCursor:
+
+    def __init__(self, tx):
+        self._tx = tx
+        # TODO: assuming here that the keys are not changed while the cursor exists.
+        # Any way to enforce it?
+        self._keys = list(sorted(tx._storage))
+        self._pos = None
+
+    def set_range(self, key):
+        pos = bisect_left(self._keys, key)
+        if pos == len(self._keys):
+            self._pos = None
+            return False
+        else:
+            self._pos = pos
+            return True
+
+    def key(self):
+        return self._keys[self._pos]
+
+    def iternext(self, keys=True, values=True):
+        return iter(self._keys[self._pos:])

tests/mock/performance_mocks.py

@@ -168,9 +168,10 @@ class NotARestApp:
     _actual_rest_apps = []
     _replaced_routes = {}
 
-    def __init__(self, this_node, *args, **kwargs):
+    def __init__(self, this_node, db_filepath, *args, **kwargs):
         self._actual_rest_app = None
         self.this_node = this_node
+        self.db_filepath = db_filepath
 
     @classmethod
     def create_with_not_a_datastore(cls, *args, **kwargs):
@@ -195,7 +196,7 @@ class NotARestApp:
 
     def actual_rest_app(self):
         if self._actual_rest_app is None:
-            self._actual_rest_app, self._datastore = make_rest_app(db_filepath=tempfile.mkdtemp(),
+            self._actual_rest_app, self._datastore = make_rest_app(db_filepath=self.db_filepath,
                                                                    this_node=self.this_node,
                                                                    domain=None)
             _new_view_functions = self._ViewFunctions(self._actual_rest_app.view_functions)

tests/unit/datastore/conftest.py

@@ -0,0 +1,57 @@
+"""
+This file is part of nucypher.
+
+nucypher is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+nucypher is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+from constant_sorrow.constants import MOCK_DB
+import lmdb
+import pytest
+import shutil
+import tempfile
+
+from nucypher.datastore import datastore
+
+
+@pytest.fixture(scope='function')
+def mock_or_real_datastore(request):
+    if request.param:
+        yield datastore.Datastore(MOCK_DB)
+    else:
+        temp_path = tempfile.mkdtemp()
+        yield datastore.Datastore(temp_path)
+        shutil.rmtree(temp_path)
+
+
+@pytest.fixture(scope='function')
+def mock_or_real_lmdb_env(request):
+    if request.param:
+        yield lmdb.open(MOCK_DB)
+    else:
+        temp_path = tempfile.mkdtemp()
+        yield lmdb.open(temp_path)
+        shutil.rmtree(temp_path)
+
+
+def pytest_generate_tests(metafunc):
+
+    if 'mock_or_real_datastore' in metafunc.fixturenames:
+        values = [False, True]
+        ids = ['real_datastore', 'mock_datastore']
+        metafunc.parametrize('mock_or_real_datastore', values, ids=ids, indirect=True)
+
+    if 'mock_or_real_lmdb_env' in metafunc.fixturenames:
+        values = [False, True]
+        ids = ['real_lmdb_env', 'mock_lmdb_env']
+        metafunc.parametrize('mock_or_real_lmdb_env', values, ids=ids, indirect=True)

tests/unit/datastore/test_datastore.py

@@ -32,13 +32,18 @@ class TestRecord(DatastoreRecord):
             decode=lambda val: datetime.fromisoformat(val.decode()))
 
 
-def test_datastore_describe():
+def test_datastore_create():
     temp_path = tempfile.mkdtemp()
     storage = datastore.Datastore(temp_path)
     assert storage.LMDB_MAP_SIZE == 1_000_000_000_000
     assert storage.db_path == temp_path
     assert storage._Datastore__db_env.path() == temp_path
 
+
+def test_datastore_describe(mock_or_real_datastore):
+
+    storage = mock_or_real_datastore
+
     #
     # Tests for `Datastore.describe`
     #
@@ -153,9 +158,9 @@ def test_datastore_describe():
         assert new_test_record.test == b'now it exists :)'
 
 
-def test_datastore_query_by():
-    temp_path = tempfile.mkdtemp()
-    storage = datastore.Datastore(temp_path)
+def test_datastore_query_by(mock_or_real_datastore):
+
+    storage = mock_or_real_datastore
 
     # Make two test record classes
     class FooRecord(DatastoreRecord):
@@ -253,8 +258,8 @@ def test_datastore_query_by():
             assert len(records) == 'this never gets executed'
 
 
-def test_datastore_record_read():
-    db_env = lmdb.open(tempfile.mkdtemp())
+def test_datastore_record_read(mock_or_real_lmdb_env):
+    db_env = mock_or_real_lmdb_env
     with db_env.begin() as db_tx:
         # Check the default attrs.
         test_rec = TestRecord(db_tx, 'testing', writeable=False)
@@ -276,9 +281,9 @@ def test_datastore_record_read():
             test_rec.test = b'should error'
 
 
-def test_datastore_record_write():
+def test_datastore_record_write(mock_or_real_lmdb_env):
     # Test writing
-    db_env = lmdb.open(tempfile.mkdtemp())
+    db_env = mock_or_real_lmdb_env
     with db_env.begin(write=True) as db_tx:
         test_rec = TestRecord(db_tx, 'testing', writeable=True)
         assert test_rec._DatastoreRecord__writeable == True
@@ -305,6 +310,9 @@ def test_datastore_record_write():
         # TODO: Mock a `DBWriteError`
 
     # Test abort
+    # Transaction context manager attempts to commit the transaction at `__exit__()`,
+    # since there were no errors caught, but it's already been aborted,
+    # so `lmdb.Error` is raised.
     with pytest.raises(lmdb.Error):
         with db_env.begin(write=True) as db_tx:
             test_rec = TestRecord(db_tx, 'testing', writeable=True)

tests/unit/datastore/test_datastore_models.py

@@ -23,9 +23,8 @@ from nucypher.datastore import datastore
 from nucypher.datastore.models import PolicyArrangement, TreasureMap, Workorder
 
 
-def test_policy_arrangement_model():
-    temp_path = tempfile.mkdtemp()
-    storage = datastore.Datastore(temp_path)
+def test_policy_arrangement_model(mock_or_real_datastore):
+    storage = mock_or_real_datastore
 
     arrangement_id_hex = 'beef'
     expiration = maya.now()
@@ -51,9 +50,8 @@ def test_policy_arrangement_model():
             should_error = policy_arrangement.arrangement_id
 
 
-def test_workorder_model():
-    temp_path = tempfile.mkdtemp()
-    storage = datastore.Datastore(temp_path)
+def test_workorder_model(mock_or_real_datastore):
+    storage = mock_or_real_datastore
     bob_keypair = keypairs.SigningKeypair(generate_keys_if_needed=True)
 
     arrangement_id_hex = 'beef'
@@ -80,9 +78,8 @@ def test_workorder_model():
             should_error = work_order.arrangement_id
 
 
-def test_treasure_map_model():
-    temp_path = tempfile.mkdtemp()
-    storage = datastore.Datastore(temp_path)
+def test_treasure_map_model(mock_or_real_datastore):
+    storage = mock_or_real_datastore
 
     hrac = 'beef'
     fake_treasure_map_data = b'My Little TreasureMap'

tests/utils/ursula.py

@@ -33,9 +33,9 @@ from nucypher.config.characters import UrsulaConfiguration
 from nucypher.crypto.powers import TransactingPower
 from nucypher.policy.collections import WorkOrder, IndisputableEvidence
 from tests.constants import (
-    MOCK_URSULA_DB_FILEPATH,
     NUMBER_OF_URSULAS_IN_DEVELOPMENT_NETWORK
 )
+from tests.mock.datastore import MOCK_DB
 from umbral import pre
 from umbral.curvebn import CurveBN
 from umbral.keys import UmbralPrivateKey
@@ -77,7 +77,7 @@ def make_federated_ursulas(ursula_config: UrsulaConfiguration,
     for port in range(starting_port, starting_port+quantity):
 
         ursula = ursula_config.produce(rest_port=port + 100,
-                                       db_filepath=MOCK_URSULA_DB_FILEPATH,
+                                       db_filepath=MOCK_DB,
                                        **ursula_overrides)
 
         federated_ursulas.add(ursula)
@@ -112,7 +112,7 @@ def make_decentralized_ursulas(ursula_config: UrsulaConfiguration,
     for port, (staker_address, worker_address) in enumerate(stakers_and_workers, start=starting_port):
         ursula = ursula_config.produce(checksum_address=staker_address,
                                        worker_address=worker_address,
-                                       db_filepath=tempfile.mkdtemp(),
+                                       db_filepath=MOCK_DB,
                                        rest_port=port + 100,
                                        # start_working_now=commit_to_next_period,  # FIXME: 2424
                                        **ursula_overrides)