873 lines
37 KiB
Python
873 lines
37 KiB
Python
# SPDX-License-Identifier: LGPL-2.1-or-later
|
|
|
|
"""
|
|
AssetManager tests.
|
|
"""
|
|
|
|
import unittest
|
|
import asyncio
|
|
from unittest.mock import Mock
|
|
import pathlib
|
|
import tempfile
|
|
from typing import Any, Mapping, List, Type, Optional, cast
|
|
from Path.Tool.assets import (
|
|
AssetManager,
|
|
FileStore,
|
|
Asset,
|
|
AssetUri,
|
|
MemoryStore,
|
|
AssetSerializer,
|
|
DummyAssetSerializer,
|
|
)
|
|
|
|
|
|
# Mock Asset class for testing
|
|
class MockAsset(Asset):
|
|
asset_type: str = "mock_asset"
|
|
|
|
def __init__(self, data: Any = None, id: str = "mock_id"):
|
|
self._data = data
|
|
self._id = id
|
|
|
|
@classmethod
|
|
def extract_dependencies(cls, data: bytes, serializer: Type[AssetSerializer]) -> List[AssetUri]:
|
|
# Mock implementation doesn't use data or format for dependencies
|
|
return []
|
|
|
|
@classmethod
|
|
def from_bytes(
|
|
cls,
|
|
data: bytes,
|
|
id: str,
|
|
dependencies: Optional[Mapping[AssetUri, Asset]],
|
|
serializer: Type[AssetSerializer],
|
|
) -> "MockAsset":
|
|
# Create instance with provided id
|
|
return cls(data, id)
|
|
|
|
def to_bytes(self, serializer: Type[AssetSerializer]) -> bytes:
|
|
return self._data
|
|
|
|
def get_id(self) -> str:
|
|
return self._id
|
|
|
|
def get_data(self) -> bytes:
|
|
"""Returns the raw data stored in the mock asset."""
|
|
return self._data
|
|
|
|
|
|
# Mock Asset class with dependencies for testing deepcopy
|
|
class MockAssetWithDeps(Asset):
|
|
asset_type: str = "mock_asset_with_deps"
|
|
|
|
def __init__(
|
|
self,
|
|
data: Any = None,
|
|
id: str = "mock_id",
|
|
dependencies: Optional[Mapping[AssetUri, Asset]] = None,
|
|
):
|
|
self._data = data
|
|
self._id = id
|
|
self._dependencies = dependencies or {}
|
|
|
|
@classmethod
|
|
def extract_dependencies(cls, data: bytes, serializer: Type[AssetSerializer]) -> List[AssetUri]:
|
|
# Assuming data is a simple JSON string like '{"deps": ["uri1", "uri2"]}'
|
|
try:
|
|
import json
|
|
|
|
data_str = data.decode("utf-8")
|
|
data_dict = json.loads(data_str)
|
|
dep_uris_str = data_dict.get("deps", [])
|
|
return [AssetUri(uri_str) for uri_str in dep_uris_str]
|
|
except Exception:
|
|
return []
|
|
|
|
@classmethod
|
|
def from_bytes(
|
|
cls,
|
|
data: bytes,
|
|
id: str,
|
|
dependencies: Optional[Mapping[AssetUri, Asset]],
|
|
serializer: Type[AssetSerializer],
|
|
) -> "MockAssetWithDeps":
|
|
# Create instance with provided id and resolved dependencies
|
|
return cls(data, id, dependencies)
|
|
|
|
def to_bytes(self, serializer: Type[AssetSerializer]) -> bytes:
|
|
# Serialize data and dependency URIs into a simple format
|
|
try:
|
|
import json
|
|
|
|
dep_uri_strs = [str(uri) for uri in self._dependencies.keys()]
|
|
data_dict = {"data": self._data.decode("utf-8"), "deps": dep_uri_strs}
|
|
return json.dumps(data_dict).encode("utf-8")
|
|
except Exception:
|
|
return self._data # Fallback if serialization fails
|
|
|
|
def get_id(self) -> str:
|
|
return self._id
|
|
|
|
def get_data(self) -> bytes:
|
|
"""Returns the raw data stored in the mock asset."""
|
|
return self._data
|
|
|
|
def get_dependencies(self) -> Mapping[AssetUri, Asset]:
|
|
"""Returns the resolved dependencies."""
|
|
return self._dependencies
|
|
|
|
|
|
class TestPathToolAssetManager(unittest.TestCase):
|
|
def test_register_store(self):
|
|
manager = AssetManager()
|
|
mock_store_local = Mock()
|
|
mock_store_local.name = "local"
|
|
mock_store_remote = Mock()
|
|
mock_store_remote.name = "remote"
|
|
|
|
manager.register_store(mock_store_local)
|
|
self.assertEqual(manager.stores["local"], mock_store_local)
|
|
|
|
manager.register_store(mock_store_remote)
|
|
self.assertEqual(manager.stores["remote"], mock_store_remote)
|
|
|
|
# Test overwriting
|
|
mock_store_local_new = Mock()
|
|
mock_store_local_new.name = "local"
|
|
manager.register_store(mock_store_local_new)
|
|
self.assertEqual(manager.stores["local"], mock_store_local_new)
|
|
|
|
def test_register_asset(self):
|
|
manager = AssetManager()
|
|
# Register the actual MockAsset class
|
|
manager.register_asset(MockAsset, DummyAssetSerializer)
|
|
self.assertEqual(manager._asset_classes[MockAsset.asset_type], MockAsset)
|
|
|
|
# Test registering a different actual Asset class
|
|
class AnotherMockAsset(Asset):
|
|
asset_type: str = "another_mock_asset"
|
|
|
|
@classmethod
|
|
def dependencies(cls, data: bytes) -> List[AssetUri]:
|
|
return []
|
|
|
|
@classmethod
|
|
def from_bytes(
|
|
cls,
|
|
data: bytes,
|
|
id: str,
|
|
dependencies: Optional[Mapping[AssetUri, Asset]],
|
|
serializer: Type[AssetSerializer],
|
|
) -> "AnotherMockAsset":
|
|
return cls()
|
|
|
|
def to_bytes(self, serializer: Type[AssetSerializer]) -> bytes:
|
|
return b""
|
|
|
|
def get_id(self) -> str:
|
|
return "another_mock_id"
|
|
|
|
manager.register_asset(AnotherMockAsset, DummyAssetSerializer)
|
|
self.assertEqual(manager._asset_classes[AnotherMockAsset.asset_type], AnotherMockAsset)
|
|
|
|
# Test overwriting
|
|
manager.register_asset(
|
|
MockAsset, DummyAssetSerializer
|
|
) # Registering again should overwrite
|
|
self.assertEqual(manager._asset_classes[MockAsset.asset_type], MockAsset)
|
|
|
|
# Test registering non-Asset class
|
|
with self.assertRaises(TypeError):
|
|
|
|
class NotAnAsset(Asset): # Inherit from Asset
|
|
pass
|
|
|
|
manager.register_asset(NotAnAsset, DummyAssetSerializer)
|
|
|
|
def test_get(self):
|
|
# Setup AssetManager with a real LocalStore and the MockAsset class
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
base_dir = pathlib.Path(tmpdir)
|
|
local_store = FileStore("local", base_dir)
|
|
manager = AssetManager()
|
|
manager.register_store(local_store)
|
|
|
|
# Register the MockAsset class
|
|
manager.register_asset(MockAsset, DummyAssetSerializer)
|
|
|
|
# Create a test asset file via AssetManager
|
|
test_data = b"test asset data"
|
|
test_uri = manager.add_raw(
|
|
asset_type=MockAsset.asset_type,
|
|
asset_id="dummy_id_get",
|
|
data=test_data,
|
|
store="local",
|
|
)
|
|
|
|
# Call AssetManager.get
|
|
retrieved_object = cast(MockAsset, manager.get(test_uri))
|
|
|
|
# Assert the retrieved object is an instance of MockAsset
|
|
self.assertIsInstance(retrieved_object, MockAsset)
|
|
self.assertEqual(retrieved_object.get_data(), test_data)
|
|
|
|
# Test error handling for non-existent URI
|
|
non_existent_uri = AssetUri.build(MockAsset.asset_type, "non_existent", "1")
|
|
with self.assertRaises(FileNotFoundError):
|
|
manager.get(non_existent_uri)
|
|
|
|
# Test error handling for no asset class registered
|
|
non_registered_uri = AssetUri.build("non_existent_type", "dummy_id", "1")
|
|
# Need to create a dummy file for the store to find
|
|
dummy_data = b"dummy"
|
|
manager.add_raw(
|
|
asset_type="non_existent_type", asset_id="dummy_id", data=dummy_data, store="local"
|
|
)
|
|
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.get(non_registered_uri)
|
|
self.assertIn("No asset class registered for URI:", str(cm.exception))
|
|
|
|
def test_delete(self):
|
|
# Setup AssetManager with a real LocalStore
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
base_dir = pathlib.Path(tmpdir)
|
|
local_store = FileStore("local", base_dir)
|
|
manager = AssetManager()
|
|
manager.register_store(local_store)
|
|
|
|
# Create a test asset file
|
|
test_data = b"test asset data to delete"
|
|
test_uri = manager.add_raw(
|
|
asset_type="temp_asset", asset_id="dummy_id_delete", data=test_data, store="local"
|
|
)
|
|
test_path = base_dir / "temp_asset" / str(test_uri.asset_id) / str(test_uri.version)
|
|
self.assertTrue(test_path.exists())
|
|
|
|
# Call AssetManager.delete
|
|
manager.delete(test_uri)
|
|
|
|
# Verify file deletion
|
|
self.assertFalse(test_path.exists())
|
|
|
|
# Test error handling for non-existent URI (should not raise error
|
|
# as LocalStore.delete handles this)
|
|
non_existent_uri = AssetUri.build(
|
|
"temp_asset", "non_existent", "1" # Keep original for logging
|
|
)
|
|
manager.delete(non_existent_uri) # Should not raise
|
|
|
|
def test_create(self):
|
|
# Setup AssetManager with LocalStore and MockAsset class
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
base_dir = pathlib.Path(tmpdir)
|
|
local_store = FileStore("local", base_dir)
|
|
manager = AssetManager()
|
|
manager.register_store(local_store)
|
|
|
|
# Register the MockAsset class
|
|
manager.register_asset(MockAsset, DummyAssetSerializer)
|
|
|
|
# Create a MockAsset instance with a specific id
|
|
test_obj = MockAsset(b"object data", id="mocked_asset_id")
|
|
|
|
# Call manager.add to create
|
|
created_uri = manager.add(test_obj, store="local")
|
|
|
|
# Assert returned URI is as expected
|
|
expected_uri = AssetUri.build(MockAsset.asset_type, "mocked_asset_id", "1")
|
|
self.assertEqual(created_uri, expected_uri)
|
|
|
|
# Verify the asset was created
|
|
retrieved_data = asyncio.run(local_store.get(created_uri))
|
|
self.assertEqual(retrieved_data, test_obj.get_data())
|
|
|
|
# Test error handling (store not found)
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.add(test_obj, store="non_existent_store")
|
|
self.assertIn("No store registered for name:", str(cm.exception))
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
local_store = MemoryStore("local")
|
|
manager = AssetManager()
|
|
manager.register_store(local_store)
|
|
|
|
# Register the MockAsset class
|
|
manager.register_asset(MockAsset, DummyAssetSerializer)
|
|
|
|
# First, create an asset
|
|
initial_data = b"initial data"
|
|
asset_id = "some_asset_id"
|
|
test_uri = manager.add_raw(MockAsset.asset_type, asset_id, initial_data, "local")
|
|
self.assertEqual(test_uri.version, "1")
|
|
|
|
# Create a MockAsset instance with the same id for update
|
|
updated_data = b"updated object data"
|
|
test_obj = MockAsset(updated_data, id=asset_id)
|
|
|
|
# Call manager.add to update
|
|
updated_uri = manager.add(test_obj, store="local")
|
|
|
|
# Assert returned URI matches the original except for version
|
|
self.assertEqual(updated_uri.asset_type, test_uri.asset_type)
|
|
self.assertEqual(updated_uri.asset_id, test_uri.asset_id)
|
|
self.assertEqual(updated_uri.version, "2")
|
|
|
|
# Verify the asset was updated
|
|
obj = cast(MockAsset, manager.get(updated_uri, store="local"))
|
|
self.assertEqual(updated_data, test_obj.get_data())
|
|
self.assertIsInstance(obj, MockAsset)
|
|
self.assertEqual(updated_data, obj.get_data())
|
|
|
|
# Test error handling (store not found)
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.add(test_obj, store="non_existent_store")
|
|
self.assertIn("No store registered for name:", str(cm.exception))
|
|
|
|
def test_create_raw(self):
|
|
# Setup AssetManager with a real MemoryStore
|
|
memory_store = MemoryStore("memory_raw")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store)
|
|
|
|
asset_type = "raw_test_type"
|
|
asset_id = "raw_test_id"
|
|
data = b"raw test data"
|
|
|
|
# Expected URI with version 1 (assuming MemoryStore uses integer versions)
|
|
expected_uri = AssetUri.build(asset_type, asset_id, "1")
|
|
|
|
# Call manager.add_raw
|
|
created_uri = manager.add_raw(
|
|
asset_type=asset_type, asset_id=asset_id, data=data, store="memory_raw"
|
|
)
|
|
|
|
# Assert returned URI is correct (check asset_type and asset_id)
|
|
self.assertEqual(created_uri.asset_type, asset_type)
|
|
self.assertEqual(created_uri.asset_id, asset_id)
|
|
self.assertEqual(created_uri, expected_uri)
|
|
|
|
# Verify data was stored using the actual created_uri
|
|
# Await the async get method using asyncio.run
|
|
retrieved_data = asyncio.run(memory_store.get(created_uri))
|
|
self.assertEqual(retrieved_data, data)
|
|
|
|
# Test error handling (store not found)
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.add_raw(
|
|
asset_type=asset_type, asset_id=asset_id, data=data, store="non_existent_store"
|
|
)
|
|
self.assertIn("No store registered for name:", str(cm.exception))
|
|
|
|
def test_get_raw(self):
|
|
# Setup AssetManager with a real MemoryStore
|
|
memory_store = MemoryStore("memory_raw_get")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store)
|
|
|
|
test_uri_str = "test_type://test_id/1"
|
|
test_uri = AssetUri(test_uri_str)
|
|
expected_data = b"retrieved raw data"
|
|
|
|
# Manually put data into the memory store
|
|
manager.add_raw("test_type", "test_id", expected_data, "memory_raw_get")
|
|
|
|
# Call manager.get_raw using the URI returned by add_raw
|
|
retrieved_data = manager.get_raw(test_uri, store="memory_raw_get")
|
|
|
|
# Assert returned data matches store's result
|
|
self.assertEqual(retrieved_data, expected_data)
|
|
|
|
# Test error handling (store not found)
|
|
non_existent_uri = AssetUri("type://id/1")
|
|
# Test error handling (asset not found in any store, including non-existent ones)
|
|
non_existent_uri = AssetUri("type://id/1")
|
|
with self.assertRaises(FileNotFoundError) as cm:
|
|
manager.get_raw(non_existent_uri, store="non_existent_store")
|
|
self.assertIn(
|
|
"Asset 'type://id/1' not found in stores '['non_existent_store']'", str(cm.exception)
|
|
)
|
|
|
|
def test_is_empty(self):
|
|
# Setup AssetManager with a real MemoryStore
|
|
memory_store = MemoryStore("memory_empty")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store)
|
|
|
|
# Test when store is empty
|
|
self.assertTrue(manager.is_empty(store="memory_empty"))
|
|
|
|
# Add an asset and test again
|
|
manager.add_raw("test_type", "test_id", b"data", "memory_empty")
|
|
self.assertFalse(manager.is_empty(store="memory_empty"))
|
|
|
|
# Test with asset type
|
|
self.assertTrue(manager.is_empty(store="memory_empty", asset_type="another_type"))
|
|
self.assertFalse(manager.is_empty(store="memory_empty", asset_type="test_type"))
|
|
|
|
# Test error handling (store not found)
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.is_empty(store="non_existent_store")
|
|
self.assertIn("No store registered for name:", str(cm.exception))
|
|
|
|
def test_count_assets(self):
|
|
# Setup AssetManager with a real MemoryStore
|
|
memory_store = MemoryStore("memory_count")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store)
|
|
|
|
# Test when store is empty
|
|
self.assertEqual(manager.count_assets(store="memory_count"), 0)
|
|
self.assertEqual(manager.count_assets(store="memory_count", asset_type="type1"), 0)
|
|
|
|
# Add assets and test counts
|
|
manager.add_raw("type1", "asset1", b"data1", "memory_count")
|
|
self.assertEqual(manager.count_assets(store="memory_count"), 1)
|
|
self.assertEqual(manager.count_assets(store="memory_count", asset_type="type1"), 1)
|
|
self.assertEqual(manager.count_assets(store="memory_count", asset_type="type2"), 0)
|
|
|
|
manager.add_raw("type2", "asset2", b"data2", "memory_count")
|
|
manager.add_raw("type1", "asset3", b"data3", "memory_count")
|
|
self.assertEqual(manager.count_assets(store="memory_count"), 3)
|
|
self.assertEqual(manager.count_assets(store="memory_count", asset_type="type1"), 2)
|
|
self.assertEqual(manager.count_assets(store="memory_count", asset_type="type2"), 1)
|
|
|
|
# Test error handling (store not found)
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.count_assets(store="non_existent_store")
|
|
self.assertIn("No store registered for name:", str(cm.exception))
|
|
|
|
def test_get_bulk(self):
|
|
# Setup AssetManager with a real MemoryStore and MockAsset class
|
|
memory_store = MemoryStore("memory_bulk")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store)
|
|
manager.register_asset(MockAsset, DummyAssetSerializer)
|
|
|
|
# Create some assets in the memory store
|
|
data1 = b"data for id1"
|
|
data2 = b"data for id2"
|
|
uri1 = manager.add_raw(MockAsset.asset_type, "id1", data1, "memory_bulk")
|
|
uri2 = manager.add_raw(MockAsset.asset_type, "id2", data2, "memory_bulk")
|
|
uri3 = AssetUri.build(MockAsset.asset_type, "non_existent", "1")
|
|
uris = [uri1, uri2, uri3]
|
|
|
|
# Call manager.get_bulk
|
|
retrieved_assets = cast(List[MockAsset], manager.get_bulk(uris, store="memory_bulk"))
|
|
|
|
# Assert the correct number of assets were returned
|
|
self.assertEqual(len(retrieved_assets), 3)
|
|
|
|
# Assert the retrieved assets are MockAsset instances with correct data
|
|
self.assertIsInstance(retrieved_assets[0], MockAsset)
|
|
self.assertEqual(retrieved_assets[0].get_data(), data1)
|
|
|
|
self.assertIsInstance(retrieved_assets[1], MockAsset)
|
|
self.assertEqual(retrieved_assets[1].get_data(), data2)
|
|
|
|
# Assert the non-existent asset is None
|
|
self.assertIsNone(retrieved_assets[2])
|
|
|
|
# Test error handling (store not found)
|
|
# Test handling of non-existent store (should skip and not raise ValueError)
|
|
# The test already asserts that the non-existent asset is None, which is the expected behavior.
|
|
manager.get_bulk(uris, store="non_existent_store")
|
|
|
|
def test_fetch(self):
|
|
# Setup AssetManager with a real MemoryStore and MockAsset class
|
|
memory_store = MemoryStore("memory_fetch")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store)
|
|
manager.register_asset(MockAsset, DummyAssetSerializer)
|
|
|
|
# Create some assets in the memory store
|
|
data1 = b"data for id1"
|
|
data2 = b"data for id2"
|
|
manager.add_raw(MockAsset.asset_type, "id1", data1, "memory_fetch")
|
|
manager.add_raw(MockAsset.asset_type, "id2", data2, "memory_fetch")
|
|
# Create an asset of a different type
|
|
manager.add_raw("another_type", "id3", b"data for id3", "memory_fetch")
|
|
AssetUri.build(MockAsset.asset_type, "non_existent", "1")
|
|
|
|
# Call manager.fetch without filters
|
|
# This should raise ValueError because uri3 has an unregistered type
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.fetch(store="memory_fetch")
|
|
self.assertIn("No asset class registered for URI:", str(cm.exception))
|
|
|
|
# Now test fetching with a registered asset type filter
|
|
# Setup a new manager and store to avoid state from previous test
|
|
memory_store_filtered = MemoryStore("memory_fetch_filtered")
|
|
manager_filtered = AssetManager()
|
|
manager_filtered.register_store(memory_store_filtered)
|
|
manager_filtered.register_asset(MockAsset, DummyAssetSerializer)
|
|
|
|
# Create assets again
|
|
manager_filtered.add_raw(MockAsset.asset_type, "id1", data1, "memory_fetch_filtered")
|
|
manager_filtered.add_raw(MockAsset.asset_type, "id2", data2, "memory_fetch_filtered")
|
|
manager_filtered.add_raw("another_type", "id3", b"data for id3", "memory_fetch_filtered")
|
|
|
|
retrieved_assets_filtered = cast(
|
|
List[MockAsset],
|
|
manager_filtered.fetch(asset_type=MockAsset.asset_type, store="memory_fetch_filtered"),
|
|
)
|
|
|
|
# Assert the correct number of assets were returned
|
|
self.assertEqual(len(retrieved_assets_filtered), 2)
|
|
|
|
# Assert the retrieved assets are MockAsset instances with correct data
|
|
self.assertIsInstance(retrieved_assets_filtered[0], MockAsset)
|
|
self.assertEqual(
|
|
retrieved_assets_filtered[0].get_data().decode("utf-8"),
|
|
data1.decode("utf-8"),
|
|
)
|
|
|
|
self.assertIsInstance(retrieved_assets_filtered[1], MockAsset)
|
|
self.assertEqual(
|
|
retrieved_assets_filtered[1].get_data().decode("utf-8"),
|
|
data2.decode("utf-8"),
|
|
)
|
|
|
|
# Test error handling (store not found)
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.fetch(store="non_existent_store")
|
|
self.assertIn("No store registered for name:", str(cm.exception))
|
|
|
|
def test_copy(self):
|
|
# Setup AssetManager with two MemoryStores and MockAsset class
|
|
memory_store_src = MemoryStore("memory_copy_src")
|
|
memory_store_dest = MemoryStore("memory_copy_dest")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store_src)
|
|
manager.register_store(memory_store_dest)
|
|
manager.register_asset(MockAsset, DummyAssetSerializer)
|
|
|
|
# Create a source asset
|
|
src_data = b"source asset data"
|
|
src_uri = manager.add_raw(MockAsset.asset_type, "source_id", src_data, "memory_copy_src")
|
|
|
|
# Test copying to a different store with default destination URI
|
|
copied_uri_default_dest = manager.copy(
|
|
src_uri, dest_store="memory_copy_dest", store="memory_copy_src"
|
|
)
|
|
self.assertEqual(copied_uri_default_dest.asset_type, src_uri.asset_type)
|
|
self.assertEqual(copied_uri_default_dest.asset_id, src_uri.asset_id)
|
|
self.assertEqual(copied_uri_default_dest.version, "1") # First version in dest
|
|
|
|
# Verify the copied asset exists in the destination store
|
|
copied_data_default_dest = manager.get_raw(
|
|
copied_uri_default_dest, store="memory_copy_dest"
|
|
)
|
|
self.assertEqual(copied_data_default_dest, src_data)
|
|
|
|
# Test copying to a different store with a specified destination URI
|
|
dest_uri_specified = AssetUri.build(MockAsset.asset_type, "specified_dest_id", "1")
|
|
copied_uri_specified_dest = manager.copy(
|
|
src_uri,
|
|
dest_store="memory_copy_dest",
|
|
store="memory_copy_src",
|
|
dest=dest_uri_specified,
|
|
)
|
|
self.assertEqual(copied_uri_specified_dest, dest_uri_specified)
|
|
|
|
# Verify the copied asset exists at the specified destination URI
|
|
copied_data_specified_dest = manager.get_raw(
|
|
copied_uri_specified_dest, store="memory_copy_dest"
|
|
)
|
|
self.assertEqual(copied_data_specified_dest, src_data)
|
|
|
|
# Test copying to the same store with a different destination URI
|
|
dest_uri_same_store = AssetUri.build(MockAsset.asset_type, "same_store_dest", "1")
|
|
copied_uri_same_store = manager.copy(
|
|
src_uri, dest_store="memory_copy_src", store="memory_copy_src", dest=dest_uri_same_store
|
|
)
|
|
self.assertEqual(copied_uri_same_store, dest_uri_same_store)
|
|
|
|
# Verify the copied asset exists in the same store at the new URI
|
|
copied_data_same_store = manager.get_raw(copied_uri_same_store, store="memory_copy_src")
|
|
self.assertEqual(copied_data_same_store, src_data)
|
|
|
|
# Test assertion for source and destination being the same
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.copy(
|
|
src_uri, dest_store="memory_copy_src", store="memory_copy_src", dest=src_uri
|
|
)
|
|
self.assertIn(
|
|
"Source and destination cannot be the same asset in the same store.",
|
|
str(cm.exception),
|
|
)
|
|
|
|
# Test overwriting existing destination (add a different asset at specified_dest_id)
|
|
overwrite_data = b"data to be overwritten"
|
|
overwrite_uri = manager.add_raw(
|
|
MockAsset.asset_type, "specified_dest_id", overwrite_data, "memory_copy_dest"
|
|
)
|
|
self.assertEqual(overwrite_uri.version, "2") # Should be version 2 now
|
|
|
|
# Copy the original src_uri to the existing destination
|
|
copied_uri_overwrite = manager.copy(
|
|
src_uri,
|
|
dest_store="memory_copy_dest",
|
|
store="memory_copy_src",
|
|
dest=dest_uri_specified,
|
|
)
|
|
# The version should be incremented again due to overwrite
|
|
self.assertEqual(copied_uri_overwrite.version, "3")
|
|
|
|
# Verify the destination now contains the source data
|
|
retrieved_overwritten_data = manager.get_raw(copied_uri_overwrite, store="memory_copy_dest")
|
|
self.assertEqual(retrieved_overwritten_data, src_data)
|
|
|
|
def test_deepcopy(self):
|
|
# Setup AssetManager with two MemoryStores and MockAssetWithDeps class
|
|
memory_store_src = MemoryStore("memory_deepcopy_src")
|
|
memory_store_dest = MemoryStore("memory_deepcopy_dest")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store_src)
|
|
manager.register_store(memory_store_dest)
|
|
manager.register_asset(MockAssetWithDeps, DummyAssetSerializer)
|
|
|
|
# Create dependency assets in the source store
|
|
dep1_data = b'{"data": "dependency 1 data", "deps": []}'
|
|
dep2_data = b'{"data": "dependency 2 data", "deps": []}'
|
|
dep1_uri = manager.add_raw(
|
|
MockAssetWithDeps.asset_type, "dep1_id", dep1_data, "memory_deepcopy_src"
|
|
)
|
|
dep2_uri = manager.add_raw(
|
|
MockAssetWithDeps.asset_type, "dep2_id", dep2_data, "memory_deepcopy_src"
|
|
)
|
|
|
|
# Create a source asset with dependencies
|
|
src_data = (
|
|
b'{"data": "source asset data", "deps": ["'
|
|
+ str(dep1_uri).encode("utf-8")
|
|
+ b'", "'
|
|
+ str(dep2_uri).encode("utf-8")
|
|
+ b'"]}'
|
|
)
|
|
src_uri = manager.add_raw(
|
|
MockAssetWithDeps.asset_type, "source_id", src_data, "memory_deepcopy_src"
|
|
)
|
|
|
|
# Test deep copying to a different store with default destination URI
|
|
copied_uri_default_dest = manager.deepcopy(
|
|
src_uri, dest_store="memory_deepcopy_dest", store="memory_deepcopy_src"
|
|
)
|
|
self.assertEqual(copied_uri_default_dest.asset_type, src_uri.asset_type)
|
|
self.assertEqual(copied_uri_default_dest.asset_id, src_uri.asset_id)
|
|
self.assertEqual(copied_uri_default_dest.version, "1") # First version in dest
|
|
|
|
# Verify the copied top-level asset exists in the destination store
|
|
copied_asset_default_dest = cast(
|
|
MockAssetWithDeps, manager.get(copied_uri_default_dest, store="memory_deepcopy_dest")
|
|
)
|
|
self.assertIsInstance(copied_asset_default_dest, MockAssetWithDeps)
|
|
# The copied asset's data should be the serialized form including dependencies
|
|
expected_data = b'{"data": "source asset data", "deps": ["mock_asset_with_deps://dep1_id/1", "mock_asset_with_deps://dep2_id/1"]}'
|
|
self.assertEqual(copied_asset_default_dest.get_data(), expected_data)
|
|
|
|
# Verify dependencies were also copied and resolved correctly
|
|
copied_deps_default_dest = copied_asset_default_dest.get_dependencies()
|
|
self.assertEqual(len(copied_deps_default_dest), 2)
|
|
self.assertIn(dep1_uri, copied_deps_default_dest)
|
|
self.assertIn(dep2_uri, copied_deps_default_dest)
|
|
|
|
copied_dep1 = cast(MockAssetWithDeps, copied_deps_default_dest[dep1_uri])
|
|
self.assertIsInstance(copied_dep1, MockAssetWithDeps)
|
|
self.assertEqual(copied_dep1.get_data(), b'{"data": "dependency 1 data", "deps": []}')
|
|
|
|
copied_dep2 = cast(MockAssetWithDeps, copied_deps_default_dest[dep2_uri])
|
|
self.assertIsInstance(copied_dep2, MockAssetWithDeps)
|
|
self.assertEqual(copied_dep2.get_data(), b'{"data": "dependency 2 data", "deps": []}')
|
|
|
|
# Test deep copying with a specified destination URI for the top-level asset
|
|
dest_uri_specified = AssetUri.build(MockAssetWithDeps.asset_type, "specified_dest_id", "1")
|
|
copied_uri_specified_dest = manager.deepcopy(
|
|
src_uri,
|
|
dest_store="memory_deepcopy_dest",
|
|
store="memory_deepcopy_src",
|
|
dest=dest_uri_specified,
|
|
)
|
|
self.assertEqual(copied_uri_specified_dest, dest_uri_specified)
|
|
|
|
# Verify the copied asset exists at the specified destination URI
|
|
copied_asset_specified_dest = cast(
|
|
MockAssetWithDeps, manager.get(copied_uri_specified_dest, store="memory_deepcopy_dest")
|
|
)
|
|
self.assertIsInstance(copied_asset_specified_dest, MockAssetWithDeps)
|
|
self.assertEqual(
|
|
copied_asset_specified_dest.get_data(),
|
|
b'{"data": "source asset data", "deps": ["mock_asset_with_deps://dep1_id/1", "mock_asset_with_deps://dep2_id/1"]}',
|
|
)
|
|
|
|
# Verify dependencies were copied and resolved correctly (their URIs should be
|
|
# in the destination store, but their asset_type and asset_id should be the same)
|
|
copied_deps_specified_dest = copied_asset_specified_dest.get_dependencies()
|
|
self.assertEqual(len(copied_deps_specified_dest), 2)
|
|
|
|
# The keys in the dependencies mapping should be the *original* URIs,
|
|
# but the values should be the *copied* dependency assets.
|
|
self.assertIn(dep1_uri, copied_deps_specified_dest)
|
|
self.assertIn(dep2_uri, copied_deps_specified_dest)
|
|
|
|
copied_dep1_specified = cast(MockAssetWithDeps, copied_deps_specified_dest[dep1_uri])
|
|
self.assertIsInstance(copied_dep1_specified, MockAssetWithDeps)
|
|
self.assertEqual(
|
|
copied_dep1_specified.get_data(), b'{"data": "dependency 1 data", "deps": []}'
|
|
)
|
|
# Check the URI of the copied dependency in the destination store
|
|
self.assertIsNotNone(
|
|
manager.get_or_none(copied_dep1_specified.get_uri(), store="memory_deepcopy_dest")
|
|
)
|
|
self.assertEqual(copied_dep1_specified.get_uri().asset_type, dep1_uri.asset_type)
|
|
self.assertEqual(copied_dep1_specified.get_uri().asset_id, dep1_uri.asset_id)
|
|
|
|
copied_dep2_specified = cast(MockAssetWithDeps, copied_deps_specified_dest[dep2_uri])
|
|
self.assertIsInstance(copied_dep2_specified, MockAssetWithDeps)
|
|
self.assertEqual(
|
|
copied_dep2_specified.get_data(), b'{"data": "dependency 2 data", "deps": []}'
|
|
)
|
|
# Check the URI of the copied dependency in the destination store
|
|
self.assertIsNotNone(
|
|
manager.get_or_none(copied_dep2_specified.get_uri(), store="memory_deepcopy_dest")
|
|
)
|
|
self.assertEqual(copied_dep2_specified.get_uri().asset_type, dep2_uri.asset_type)
|
|
self.assertEqual(copied_dep2_specified.get_uri().asset_id, dep2_uri.asset_id)
|
|
|
|
# Test handling of existing dependencies in the destination store (should be skipped)
|
|
# Add a dependency with the same URI as dep1_uri to the destination store
|
|
existing_dep1_data = b'{"data": "existing dependency 1 data", "deps": []}'
|
|
existing_dep1_uri_in_dest = manager.add_raw(
|
|
dep1_uri.asset_type, dep1_uri.asset_id, existing_dep1_data, "memory_deepcopy_dest"
|
|
)
|
|
self.assertEqual(existing_dep1_uri_in_dest.version, "2") # Should be version 2 now
|
|
|
|
# Deep copy the source asset again
|
|
copied_uri_existing_dep = manager.deepcopy(
|
|
src_uri, dest_store="memory_deepcopy_dest", store="memory_deepcopy_src"
|
|
)
|
|
# The top-level asset should be overwritten, incrementing its version
|
|
self.assertEqual(copied_uri_existing_dep.version, "2")
|
|
|
|
# Verify the top-level asset was overwritten
|
|
copied_asset_existing_dep = cast(
|
|
MockAssetWithDeps, manager.get(copied_uri_existing_dep, store="memory_deepcopy_dest")
|
|
)
|
|
self.assertIsInstance(copied_asset_existing_dep, MockAssetWithDeps)
|
|
self.assertEqual(
|
|
copied_asset_existing_dep.get_data(),
|
|
b'{"data": "source asset data", "deps": ["mock_asset_with_deps://dep1_id/1", "mock_asset_with_deps://dep2_id/1"]}',
|
|
)
|
|
|
|
# Verify that the existing dependency was *not* overwritten
|
|
retrieved_existing_dep1 = manager.get_raw(
|
|
existing_dep1_uri_in_dest, store="memory_deepcopy_dest"
|
|
)
|
|
self.assertEqual(retrieved_existing_dep1, existing_dep1_data)
|
|
|
|
# Verify the dependencies in the copied asset still point to the correct
|
|
# (existing) dependency in the destination store.
|
|
copied_deps_existing_dep = copied_asset_existing_dep.get_dependencies()
|
|
self.assertEqual(len(copied_deps_existing_dep), 2)
|
|
self.assertIn(dep1_uri, copied_deps_existing_dep)
|
|
self.assertIn(dep2_uri, copied_deps_existing_dep)
|
|
|
|
copied_dep1_existing = cast(MockAssetWithDeps, copied_deps_existing_dep[dep1_uri])
|
|
self.assertIsInstance(copied_dep1_existing, MockAssetWithDeps)
|
|
self.assertEqual(
|
|
copied_dep1_existing.get_data(), b'{"data": "dependency 1 data", "deps": []}'
|
|
) # Should be the original data from source
|
|
|
|
copied_dep2_existing = cast(MockAssetWithDeps, copied_deps_existing_dep[dep2_uri])
|
|
self.assertIsInstance(copied_dep2_existing, MockAssetWithDeps)
|
|
self.assertEqual(
|
|
copied_dep2_existing.get_data(), b'{"data": "dependency 2 data", "deps": []}'
|
|
) # Should be the newly copied raw data
|
|
|
|
# Test handling of existing top-level asset in the destination store (should be overwritten)
|
|
# This was implicitly tested in the previous step where the top-level asset's
|
|
# version was incremented. Let's add a more explicit test.
|
|
overwrite_src_data = b'{"data": "overwrite source data", "deps": []}'
|
|
overwrite_src_uri = manager.add_raw(
|
|
MockAssetWithDeps.asset_type,
|
|
"overwrite_source_id",
|
|
overwrite_src_data,
|
|
"memory_deepcopy_src",
|
|
)
|
|
|
|
# Add an asset to the destination store with the same URI as overwrite_src_uri
|
|
existing_dest_data = b'{"data": "existing destination data", "deps": []}'
|
|
existing_dest_uri = manager.add_raw(
|
|
overwrite_src_uri.asset_type,
|
|
overwrite_src_uri.asset_id,
|
|
existing_dest_data,
|
|
"memory_deepcopy_dest",
|
|
)
|
|
self.assertEqual(existing_dest_uri.version, "1")
|
|
|
|
# Deep copy overwrite_src_uri to the existing destination URI
|
|
copied_uri_overwrite_top = manager.deepcopy(
|
|
overwrite_src_uri,
|
|
dest_store="memory_deepcopy_dest",
|
|
store="memory_deepcopy_src",
|
|
dest=existing_dest_uri,
|
|
)
|
|
# The version should be incremented
|
|
self.assertEqual(copied_uri_overwrite_top.version, "2")
|
|
|
|
# Verify the destination now contains the source data
|
|
retrieved_overwritten_top = manager.get_raw(
|
|
copied_uri_overwrite_top, store="memory_deepcopy_dest"
|
|
)
|
|
# Need to parse the data to get the actual content
|
|
import json
|
|
|
|
retrieved_data_dict = json.loads(retrieved_overwritten_top.decode("utf-8"))
|
|
self.assertEqual(retrieved_data_dict.get("data"), b"overwrite source data".decode("utf-8"))
|
|
|
|
# Test error handling for non-existent source asset
|
|
non_existent_src_uri = AssetUri.build(MockAssetWithDeps.asset_type, "non_existent_src", "1")
|
|
with self.assertRaises(FileNotFoundError) as cm:
|
|
manager.deepcopy(
|
|
non_existent_src_uri, dest_store="memory_deepcopy_dest", store="memory_deepcopy_src"
|
|
)
|
|
self.assertIn("Source asset", str(cm.exception))
|
|
self.assertIn("not found", str(cm.exception))
|
|
|
|
# Test error handling for non-existent source store
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.deepcopy(src_uri, dest_store="memory_deepcopy_dest", store="non_existent_store")
|
|
self.assertIn("Source store", str(cm.exception))
|
|
self.assertIn("not registered", str(cm.exception))
|
|
|
|
# Test error handling for non-existent destination store
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.deepcopy(src_uri, dest_store="non_existent_store", store="memory_deepcopy_src")
|
|
self.assertIn("Destination store", str(cm.exception))
|
|
self.assertIn("not registered", str(cm.exception))
|
|
|
|
def test_exists(self):
|
|
# Setup AssetManager with a MemoryStore
|
|
memory_store = MemoryStore("memory_exists")
|
|
manager = AssetManager()
|
|
manager.register_store(memory_store)
|
|
|
|
# Create an asset
|
|
test_uri = manager.add_raw("test_type", "test_id", b"data", "memory_exists")
|
|
|
|
# Test exists for an existing asset
|
|
self.assertTrue(manager.exists(test_uri, store="memory_exists"))
|
|
|
|
# Test exists for a non-existent asset
|
|
non_existent_uri = AssetUri.build("test_type", "non_existent_id", "1")
|
|
self.assertFalse(manager.exists(non_existent_uri, store="memory_exists"))
|
|
|
|
# Test exists for a non-existent store (should raise ValueError)
|
|
with self.assertRaises(ValueError) as cm:
|
|
manager.exists(test_uri, store="non_existent_store")
|
|
self.assertIn("No store registered for name:", str(cm.exception))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|