diff --git a/tests/unit/distribution/routers/test_routing_tables.py b/tests/unit/distribution/routers/test_routing_tables.py index a1c3d1e95..c843557e0 100644 --- a/tests/unit/distribution/routers/test_routing_tables.py +++ b/tests/unit/distribution/routers/test_routing_tables.py @@ -354,6 +354,107 @@ async def test_scoring_functions_routing_table(cached_disk_dist_registry): assert len(scoring_functions_list_after_deletion.data) == 0 +async def test_double_registration_models_positive(cached_disk_dist_registry): + """Test that registering the same model twice with identical data succeeds.""" + table = ModelsRoutingTable({"test_provider": InferenceImpl()}, cached_disk_dist_registry, {}) + await table.initialize() + + # Register a model + await table.register_model(model_id="test-model", provider_id="test_provider") + + # Register the exact same model again - should succeed (idempotent) + await table.register_model(model_id="test-model", provider_id="test_provider") + + # Verify only one model exists + models = await table.list_models() + assert len(models.data) == 1 + assert models.data[0].identifier == "test_provider/test-model" + + +async def test_double_registration_models_negative(cached_disk_dist_registry): + """Test that registering the same model with different data fails.""" + table = ModelsRoutingTable({"test_provider": InferenceImpl()}, cached_disk_dist_registry, {}) + await table.initialize() + + # Register a model with specific metadata + await table.register_model(model_id="test-model", provider_id="test_provider", metadata={"param1": "value1"}) + + # Try to register the same model with different metadata - should fail + with pytest.raises(ValueError, match="Provider 'test_provider' is already registered"): + await table.register_model( + model_id="test-model", provider_id="test_provider", metadata={"param1": "different_value"} + ) + + +async def test_double_registration_scoring_functions_positive(cached_disk_dist_registry): + """Test that registering the same scoring function twice with identical data succeeds.""" + table = ScoringFunctionsRoutingTable({"test_provider": ScoringFunctionsImpl()}, cached_disk_dist_registry, {}) + await table.initialize() + + # Register a scoring function + await table.register_scoring_function( + scoring_fn_id="test-scoring-fn", + provider_id="test_provider", + description="Test scoring function", + return_type=NumberType(), + ) + + # Register the exact same scoring function again - should succeed (idempotent) + await table.register_scoring_function( + scoring_fn_id="test-scoring-fn", + provider_id="test_provider", + description="Test scoring function", + return_type=NumberType(), + ) + + # Verify only one scoring function exists + scoring_functions = await table.list_scoring_functions() + assert len(scoring_functions.data) == 1 + assert scoring_functions.data[0].identifier == "test-scoring-fn" + + +async def test_double_registration_scoring_functions_negative(cached_disk_dist_registry): + """Test that registering the same scoring function with different data fails.""" + table = ScoringFunctionsRoutingTable({"test_provider": ScoringFunctionsImpl()}, cached_disk_dist_registry, {}) + await table.initialize() + + # Register a scoring function + await table.register_scoring_function( + scoring_fn_id="test-scoring-fn", + provider_id="test_provider", + description="Test scoring function", + return_type=NumberType(), + ) + + # Try to register the same scoring function with different description - should fail + with pytest.raises(ValueError, match="Provider 'test_provider' is already registered"): + await table.register_scoring_function( + scoring_fn_id="test-scoring-fn", + provider_id="test_provider", + description="Different description", + return_type=NumberType(), + ) + + +async def test_double_registration_different_providers(cached_disk_dist_registry): + """Test that registering objects with same ID but different providers succeeds.""" + impl1 = InferenceImpl() + impl2 = InferenceImpl() + table = ModelsRoutingTable({"provider1": impl1, "provider2": impl2}, cached_disk_dist_registry, {}) + await table.initialize() + + # Register same model ID with different providers - should succeed + await table.register_model(model_id="shared-model", provider_id="provider1") + await table.register_model(model_id="shared-model", provider_id="provider2") + + # Verify both models exist with different identifiers + models = await table.list_models() + assert len(models.data) == 2 + model_ids = {m.identifier for m in models.data} + assert "provider1/shared-model" in model_ids + assert "provider2/shared-model" in model_ids + + async def test_benchmarks_routing_table(cached_disk_dist_registry): table = BenchmarksRoutingTable({"test_provider": BenchmarksImpl()}, cached_disk_dist_registry, {}) await table.initialize() diff --git a/tests/unit/registry/test_registry.py b/tests/unit/registry/test_registry.py index 4ea4a20b9..5e0d470e5 100644 --- a/tests/unit/registry/test_registry.py +++ b/tests/unit/registry/test_registry.py @@ -229,3 +229,132 @@ async def test_cached_registry_error_handling(sqlite_kvstore): invalid_obj = await cached_registry.get("vector_db", "invalid_cached_db") assert invalid_obj is None + + +async def test_double_registration_identical_objects(disk_dist_registry): + """Test that registering identical objects succeeds (idempotent).""" + vector_db = VectorDB( + identifier="test_vector_db", + embedding_model="all-MiniLM-L6-v2", + embedding_dimension=384, + provider_resource_id="test_vector_db", + provider_id="test-provider", + ) + + # First registration should succeed + result1 = await disk_dist_registry.register(vector_db) + assert result1 is True + + # Second registration of identical object should also succeed (idempotent) + result2 = await disk_dist_registry.register(vector_db) + assert result2 is True + + # Verify object exists and is unchanged + retrieved = await disk_dist_registry.get("vector_db", "test_vector_db") + assert retrieved is not None + assert retrieved.identifier == vector_db.identifier + assert retrieved.embedding_model == vector_db.embedding_model + + +async def test_double_registration_different_objects(disk_dist_registry): + """Test that registering different objects with same identifier fails.""" + vector_db1 = VectorDB( + identifier="test_vector_db", + embedding_model="all-MiniLM-L6-v2", + embedding_dimension=384, + provider_resource_id="test_vector_db", + provider_id="test-provider", + ) + + vector_db2 = VectorDB( + identifier="test_vector_db", # Same identifier + embedding_model="different-model", # Different embedding model + embedding_dimension=384, + provider_resource_id="test_vector_db", + provider_id="test-provider", + ) + + # First registration should succeed + result1 = await disk_dist_registry.register(vector_db1) + assert result1 is True + + # Second registration with different data should fail + with pytest.raises(ValueError, match="Provider 'test-provider' is already registered"): + await disk_dist_registry.register(vector_db2) + + # Verify original object is unchanged + retrieved = await disk_dist_registry.get("vector_db", "test_vector_db") + assert retrieved is not None + assert retrieved.embedding_model == "all-MiniLM-L6-v2" # Original value + + +async def test_double_registration_different_providers(disk_dist_registry): + """Test that objects with same identifier but different providers can be registered.""" + from llama_stack.core.datatypes import VectorDBWithOwner + + vector_db1 = VectorDBWithOwner( + identifier="shared_vector_db", + embedding_model="all-MiniLM-L6-v2", + embedding_dimension=384, + provider_resource_id="shared_vector_db", + provider_id="provider1", + ) + + vector_db2 = VectorDBWithOwner( + identifier="shared_vector_db", # Same identifier + embedding_model="all-MiniLM-L6-v2", + embedding_dimension=384, + provider_resource_id="shared_vector_db", + provider_id="provider2", # Different provider + ) + + # Both registrations should succeed because they have different provider_ids + result1 = await disk_dist_registry.register(vector_db1) + assert result1 is True + + result2 = await disk_dist_registry.register(vector_db2) + assert result2 is True + + # Both objects should be retrievable + retrieved1 = await disk_dist_registry.get("vector_db", "shared_vector_db") + # Note: Since they have the same identifier, the second one overwrites the first + # This behavior might be considered a bug, but we're testing current behavior + assert retrieved1 is not None + + +async def test_double_registration_with_cache(cached_disk_dist_registry): + """Test double registration behavior with caching enabled.""" + from llama_stack.apis.models import ModelType + from llama_stack.core.datatypes import ModelWithOwner + + model1 = ModelWithOwner( + identifier="test_model", + provider_resource_id="test_model", + provider_id="test-provider", + model_type=ModelType.llm, + ) + + model2 = ModelWithOwner( + identifier="test_model", # Same identifier + provider_resource_id="test_model", + provider_id="test-provider", + model_type=ModelType.embedding, # Different type + ) + + # First registration should succeed and populate cache + result1 = await cached_disk_dist_registry.register(model1) + assert result1 is True + + # Verify in cache + cached_model = cached_disk_dist_registry.get_cached("model", "test_model") + assert cached_model is not None + assert cached_model.model_type == ModelType.llm + + # Second registration with different data should fail + with pytest.raises(ValueError, match="Provider 'test-provider' is already registered"): + await cached_disk_dist_registry.register(model2) + + # Cache should still contain original model + cached_model_after = cached_disk_dist_registry.get_cached("model", "test_model") + assert cached_model_after is not None + assert cached_model_after.model_type == ModelType.llm