diff --git a/donkeycar/parts/datastore_v2.py b/donkeycar/parts/datastore_v2.py index 3de4a8457..8e0e01880 100644 --- a/donkeycar/parts/datastore_v2.py +++ b/donkeycar/parts/datastore_v2.py @@ -292,6 +292,10 @@ def delete_record(self, record_index): self.deleted_indexes.add(record_index) self._update_catalog_metadata(update=True) + def update_metadata(self, metadata): + self.metadata = {**self.metadata, **metadata} + self._write_contents() + def restore_record(self, record_index): # Does not actually delete the record, but marks it as deleted. self.deleted_indexes.discard(record_index) diff --git a/donkeycar/tests/test_datastore_v2.py b/donkeycar/tests/test_datastore_v2.py index 1edad5ac6..0ccbd0191 100644 --- a/donkeycar/tests/test_datastore_v2.py +++ b/donkeycar/tests/test_datastore_v2.py @@ -4,28 +4,35 @@ import time import unittest from pathlib import Path +import json -from donkeycar.parts.datastore_v2 import Manifest +from donkeycar.parts.datastore_v2 import Manifest, ManifestIterator class TestDatastore(unittest.TestCase): def setUp(self): self._path = tempfile.mkdtemp() + print(self._path) def test_basic_datastore_operations(self): # 2 records per catalog entry in the manifest - manifest = Manifest(self._path, max_len=2) - count = 10 - for i in range(count): - manifest.write_record(self._newRecord()) + max_len = 2 + record_count = 10 + metadata = dict() + # metadata + metadata = [("m0", "v0"), ("m1", "v1")] + + manifest = self._newManifest(max_len, record_count, metadata) read_records = 0 for entry in manifest: print('Entry %s' % (entry)) read_records += 1 - self.assertEqual(count, read_records) + self.assertEqual(record_count, read_records) + assert len(manifest.catalog_paths) == 5 + assert len(manifest.metadata) == 2 def test_deletion(self): manifest = Manifest(self._path, max_len=2) @@ -44,29 +51,23 @@ def test_deletion(self): self.assertEqual((count - deleted), read_records) - - def test_memory_mapped_read(self): - manifest = Manifest(self._path, max_len=2) - for i in range(10): - manifest.write_record(self._newRecord()) - manifest.close() - - manifest_2 = Manifest(self._path, read_only=True) - read_records = 0 - for _ in manifest_2: - read_records += 1 - - manifest_2.close() - - self.assertEqual(10, read_records) - - def tearDown(self): shutil.rmtree(self._path) + pass def _newRecord(self): - record = {'at' : time.time()} + record = {'at': time.time()} return record + + def _newManifest(self, max_len, record_count, metadata={}): + print(f"metadata = {metadata}") + + manifest = Manifest(self._path, metadata=metadata, max_len=max_len) + + for i in range(record_count): + manifest.write_record(self._newRecord()) + + return manifest if __name__ == '__main__': unittest.main() diff --git a/donkeycar/tests/test_tub_v2.py b/donkeycar/tests/test_tub_v2.py index e693633d3..9c265eed2 100644 --- a/donkeycar/tests/test_tub_v2.py +++ b/donkeycar/tests/test_tub_v2.py @@ -11,7 +11,8 @@ def setUp(self): self._path = tempfile.mkdtemp() inputs = ['input'] types = ['int'] - self.tub = Tub(self._path, inputs, types) + metadata = [("meta1", "metavalue1")] + self.tub = Tub(self._path, inputs, types, metadata) def test_basic_tub_operations(self): entries = list(self.tub) @@ -34,6 +35,18 @@ def test_basic_tub_operations(self): self.assertEqual(count, (write_count - len(delete_indexes))) self.assertEqual(len(self.tub), (write_count - len(delete_indexes))) + # Note that self.tub.metadata is an array of tuples + assert ("meta1", "metavalue1") in self.tub.metadata + + # Note that self.tub.manifest.metadata is a dict + assert self.tub.manifest.metadata['meta1'] == "metavalue1" + + def test_empty_tub(self): + assert len(self.tub) == 0 + tub_iter = iter(self.tub) + + self.assertRaises(StopIteration, next, tub_iter) + def tearDown(self): shutil.rmtree(self._path)