diff --git a/tests/integration/test_persistence_integration.py b/tests/integration/test_persistence_integration.py new file mode 100644 index 000000000..128a1ce9c --- /dev/null +++ b/tests/integration/test_persistence_integration.py @@ -0,0 +1,67 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import yaml + +from llama_stack.core.datatypes import StackRunConfig +from llama_stack.core.persistence_resolver import ( + resolve_inference_store_config, + resolve_metadata_store_config, +) +from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig +from llama_stack.providers.utils.sqlstore.sqlstore import ( + PostgresSqlStoreConfig, + SqliteSqlStoreConfig, +) + + +def test_starter_distribution_config_loads_and_resolves(): + """Integration: Actual starter config should parse and resolve all stores.""" + with open("llama_stack/distributions/starter/run.yaml") as f: + config_dict = yaml.safe_load(f) + + config = StackRunConfig(**config_dict) + + # Config should have persistence with default backend + assert config.persistence is not None + assert "default" in config.persistence.backends + assert isinstance(config.persistence.backends["default"], SqliteSqlStoreConfig) + + # Stores should reference the default backend + assert config.persistence.stores is not None + assert config.persistence.stores.metadata.backend == "default" + assert config.persistence.stores.inference.backend == "default" + + # Resolution should work + metadata_store = resolve_metadata_store_config(config.persistence, "starter") + assert isinstance(metadata_store, SqliteKVStoreConfig) + + sql_config, max_queue, num_writers = resolve_inference_store_config(config.persistence) + assert isinstance(sql_config, SqliteSqlStoreConfig) + assert max_queue > 0 + assert num_writers > 0 + + +def test_postgres_demo_distribution_config_loads(): + """Integration: Postgres demo should use Postgres backend for all stores.""" + with open("llama_stack/distributions/postgres-demo/run.yaml") as f: + config_dict = yaml.safe_load(f) + + config = StackRunConfig(**config_dict) + + # Should have postgres backend + assert config.persistence is not None + assert "default" in config.persistence.backends + assert isinstance(config.persistence.backends["default"], PostgresSqlStoreConfig) + + # Both stores use same postgres backend + assert config.persistence.stores.metadata.backend == "default" + assert config.persistence.stores.inference.backend == "default" + + # Resolution returns postgres config + sql_config, _, _ = resolve_inference_store_config(config.persistence) + assert isinstance(sql_config, PostgresSqlStoreConfig) + assert sql_config.host == "${env.POSTGRES_HOST:=localhost}" diff --git a/tests/unit/core/test_persistence_config.py b/tests/unit/core/test_persistence_config.py new file mode 100644 index 000000000..279f5ed15 --- /dev/null +++ b/tests/unit/core/test_persistence_config.py @@ -0,0 +1,84 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import pytest +from pydantic import ValidationError + +from llama_stack.core.datatypes import ( + InferenceStoreReference, + PersistenceConfig, + StoreReference, + StoresConfig, +) +from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig +from llama_stack.providers.utils.sqlstore.sqlstore import ( + PostgresSqlStoreConfig, + SqliteSqlStoreConfig, +) + + +def test_backend_reference_validation_catches_missing_backend(): + """Critical: Catch user typos in backend references before runtime.""" + with pytest.raises(ValidationError, match="not defined in persistence.backends"): + PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/store.db"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="typo_backend"), # User typo + ), + ) + + +def test_backend_reference_validation_accepts_valid_config(): + """Valid config should parse without errors.""" + config = PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/store.db"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="default"), + inference=InferenceStoreReference(backend="default"), + ), + ) + assert config.stores.metadata.backend == "default" + assert config.stores.inference.backend == "default" + + +def test_multiple_stores_can_share_same_backend(): + """Core use case: metadata and inference both use 'default' backend.""" + config = PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/shared.db"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="default", namespace="metadata"), + inference=InferenceStoreReference(backend="default"), + conversations=StoreReference(backend="default"), + ), + ) + # All reference the same backend + assert config.stores.metadata.backend == "default" + assert config.stores.inference.backend == "default" + assert config.stores.conversations.backend == "default" + + +def test_mixed_backend_types_allowed(): + """Should support KVStore and SqlStore backends simultaneously.""" + config = PersistenceConfig( + backends={ + "kvstore": SqliteKVStoreConfig(db_path="/tmp/kv.db"), + "sqlstore": PostgresSqlStoreConfig( + user="test", password="test", host="localhost", db="test" + ), + }, + stores=StoresConfig( + metadata=StoreReference(backend="kvstore"), + inference=InferenceStoreReference(backend="sqlstore"), + ), + ) + assert isinstance(config.backends["kvstore"], SqliteKVStoreConfig) + assert isinstance(config.backends["sqlstore"], PostgresSqlStoreConfig) diff --git a/tests/unit/core/test_persistence_resolver.py b/tests/unit/core/test_persistence_resolver.py new file mode 100644 index 000000000..24bd5ccbf --- /dev/null +++ b/tests/unit/core/test_persistence_resolver.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import pytest + +from llama_stack.core.datatypes import ( + InferenceStoreReference, + PersistenceConfig, + StoreReference, + StoresConfig, +) +from llama_stack.core.persistence_resolver import ( + resolve_backend, + resolve_inference_store_config, +) +from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig +from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig + + +def test_resolver_applies_namespace_to_kvstore(): + """Critical: Namespace overlay must work for KVStore isolation.""" + persistence = PersistenceConfig( + backends={ + "default": SqliteKVStoreConfig(db_path="/tmp/store.db"), + }, + stores=StoresConfig( + metadata=StoreReference(backend="default", namespace="meta"), + ), + ) + + resolved = resolve_backend( + persistence=persistence, + store_ref=persistence.stores.metadata, + default_factory=lambda: SqliteKVStoreConfig(db_path="/tmp/default.db"), + store_name="metadata", + ) + + # Backend config cloned with namespace applied + assert resolved.db_path == "/tmp/store.db" + assert resolved.namespace == "meta" + + +def test_resolver_does_not_apply_namespace_to_sqlstore(): + """SqlStore backends should not get namespace field.""" + persistence = PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/store.db"), + }, + stores=StoresConfig( + inference=InferenceStoreReference(backend="default"), + ), + ) + + sql_config, _, _ = resolve_inference_store_config(persistence) + + # SqlStore returned as-is, no namespace attribute + assert sql_config.db_path == "/tmp/store.db" + assert not hasattr(sql_config, "namespace") + + +def test_resolver_rejects_kvstore_for_inference(): + """Type safety: inference requires SqlStore, should fail on KVStore.""" + persistence = PersistenceConfig( + backends={ + "default": SqliteKVStoreConfig(db_path="/tmp/kv.db"), # Wrong type + }, + stores=StoresConfig( + inference=InferenceStoreReference(backend="default"), + ), + ) + + with pytest.raises(ValueError, match="requires SqlStore backend"): + resolve_inference_store_config(persistence) + + +def test_resolver_preserves_queue_params(): + """Inference store should preserve queue tuning parameters.""" + persistence = PersistenceConfig( + backends={ + "default": SqliteSqlStoreConfig(db_path="/tmp/store.db"), + }, + stores=StoresConfig( + inference=InferenceStoreReference( + backend="default", + max_write_queue_size=5000, + num_writers=2, + ), + ), + ) + + _, max_queue, num_writers = resolve_inference_store_config(persistence) + + assert max_queue == 5000 + assert num_writers == 2 + + +def test_resolver_uses_defaults_when_no_persistence_config(): + """Graceful fallback when persistence not configured.""" + sql_config, max_queue, num_writers = resolve_inference_store_config(None) + + # Should return sensible defaults + assert isinstance(sql_config, SqliteSqlStoreConfig) + assert max_queue == 10000 + assert num_writers == 4