diff --git a/tests/test_middlewares.py b/tests/test_middlewares.py index 81a21195..d9734730 100644 --- a/tests/test_middlewares.py +++ b/tests/test_middlewares.py @@ -1,5 +1,7 @@ import os +import pytest + from tinydb import TinyDB from tinydb.middlewares import CachingMiddleware from tinydb.storages import MemoryStorage, JSONStorage @@ -98,6 +100,19 @@ def test_caching_json_write(tmpdir): # Assert JSON file has been closed assert db._storage._handle.closed + +def test_caching_rejects_use_after_close(tmpdir): + path = str(tmpdir.join('closed.db')) + db = TinyDB(path, storage=CachingMiddleware(JSONStorage)) + db.insert({'key': 'value'}) + db.close() + + with pytest.raises(ValueError, match='closed'): + db.insert({'key': 'again'}) + + with pytest.raises(ValueError, match='closed'): + db.all() + del db # Reopen database diff --git a/tinydb/middlewares.py b/tinydb/middlewares.py index 7973012a..ba6a6c22 100644 --- a/tinydb/middlewares.py +++ b/tinydb/middlewares.py @@ -92,8 +92,12 @@ def __init__(self, storage_cls): # Prepare the cache self.cache = None self._cache_modified_count = 0 + self._closed = False def read(self): + if self._closed: + raise ValueError('I/O operation on closed storage') + if self.cache is None: # Empty cache: read from the storage self.cache = self.storage.read() @@ -102,6 +106,9 @@ def read(self): return self.cache def write(self, data): + if self._closed: + raise ValueError('I/O operation on closed storage') + # Store data in cache self.cache = data self._cache_modified_count += 1 @@ -114,14 +121,22 @@ def flush(self): """ Flush all unwritten data to disk. """ + if self._closed: + raise ValueError('I/O operation on closed storage') + if self._cache_modified_count > 0: # Force-flush the cache by writing the data to the storage self.storage.write(self.cache) self._cache_modified_count = 0 def close(self): + if self._closed: + return + # Flush potentially unwritten data self.flush() # Let the storage clean up too self.storage.close() + self._closed = True + self.cache = None