phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-14 09:06:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

chore: Add fixtures to conftest.py (#2067)

Browse source 
Add fixtures for SqliteKVStore, DiskDistributionRegistry and
CachedDiskDistributionRegistry. And use them in tests that had all been
duplicating similar setups.

## Test Plan
unit tests continue to run

Signed-off-by: Derek Higgins <derekh@redhat.com>
This commit is contained in:
Derek Higgins 2025-05-06 12:57:48 +01:00 • committed by GitHub
parent 4597145011
commit 2e807b38cc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 122 additions and 167 deletions
Download patch file Download diff file Expand all files Collapse all files

57
tests/unit/registry/test_registry_acl.py
View file

@ -4,9 +4,6 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import os
import shutil
import tempfile
import pytest
@ -14,30 +11,10 @@ from llama_stack.apis.models import ModelType
from llama_stack.distribution.datatypes import ModelWithACL
from llama_stack.distribution.server.auth_providers import AccessAttributes
from llama_stack.distribution.store.registry import CachedDiskDistributionRegistry
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
@pytest.fixture(scope="function")
async def kvstore():
temp_dir = tempfile.mkdtemp()
db_path = os.path.join(temp_dir, "test_registry_acl.db")
kvstore_config = SqliteKVStoreConfig(db_path=db_path)
kvstore = SqliteKVStoreImpl(kvstore_config)
await kvstore.initialize()
yield kvstore
shutil.rmtree(temp_dir)
@pytest.fixture(scope="function")
async def registry(kvstore):
registry = CachedDiskDistributionRegistry(kvstore)
await registry.initialize()
return registry
@pytest.mark.asyncio
async def test_registry_cache_with_acl(registry):
async def test_registry_cache_with_acl(cached_disk_dist_registry):
model = ModelWithACL(
identifier="model-acl",
provider_id="test-provider",
@ -46,30 +23,30 @@ async def test_registry_cache_with_acl(registry):
access_attributes=AccessAttributes(roles=["admin"], teams=["ai-team"]),
)
success = await registry.register(model)
success = await cached_disk_dist_registry.register(model)
assert success
cached_model = registry.get_cached("model", "model-acl")
cached_model = cached_disk_dist_registry.get_cached("model", "model-acl")
assert cached_model is not None
assert cached_model.identifier == "model-acl"
assert cached_model.access_attributes.roles == ["admin"]
assert cached_model.access_attributes.teams == ["ai-team"]
fetched_model = await registry.get("model", "model-acl")
fetched_model = await cached_disk_dist_registry.get("model", "model-acl")
assert fetched_model is not None
assert fetched_model.identifier == "model-acl"
assert fetched_model.access_attributes.roles == ["admin"]
model.access_attributes = AccessAttributes(roles=["admin", "user"], projects=["project-x"])
await registry.update(model)
await cached_disk_dist_registry.update(model)
updated_cached = registry.get_cached("model", "model-acl")
updated_cached = cached_disk_dist_registry.get_cached("model", "model-acl")
assert updated_cached is not None
assert updated_cached.access_attributes.roles == ["admin", "user"]
assert updated_cached.access_attributes.projects == ["project-x"]
assert updated_cached.access_attributes.teams is None
new_registry = CachedDiskDistributionRegistry(registry.kvstore)
new_registry = CachedDiskDistributionRegistry(cached_disk_dist_registry.kvstore)
await new_registry.initialize()
new_model = await new_registry.get("model", "model-acl")
@ -81,7 +58,7 @@ async def test_registry_cache_with_acl(registry):
@pytest.mark.asyncio
async def test_registry_empty_acl(registry):
async def test_registry_empty_acl(cached_disk_dist_registry):
model = ModelWithACL(
identifier="model-empty-acl",
provider_id="test-provider",
@ -90,9 +67,9 @@ async def test_registry_empty_acl(registry):
access_attributes=AccessAttributes(),
)
await registry.register(model)
await cached_disk_dist_registry.register(model)
cached_model = registry.get_cached("model", "model-empty-acl")
cached_model = cached_disk_dist_registry.get_cached("model", "model-empty-acl")
assert cached_model is not None
assert cached_model.access_attributes is not None
assert cached_model.access_attributes.roles is None
@ -100,7 +77,7 @@ async def test_registry_empty_acl(registry):
assert cached_model.access_attributes.projects is None
assert cached_model.access_attributes.namespaces is None
all_models = await registry.get_all()
all_models = await cached_disk_dist_registry.get_all()
assert len(all_models) == 1
model = ModelWithACL(
@ -110,18 +87,18 @@ async def test_registry_empty_acl(registry):
model_type=ModelType.llm,
)
await registry.register(model)
await cached_disk_dist_registry.register(model)
cached_model = registry.get_cached("model", "model-no-acl")
cached_model = cached_disk_dist_registry.get_cached("model", "model-no-acl")
assert cached_model is not None
assert cached_model.access_attributes is None
all_models = await registry.get_all()
all_models = await cached_disk_dist_registry.get_all()
assert len(all_models) == 2
@pytest.mark.asyncio
async def test_registry_serialization(registry):
async def test_registry_serialization(cached_disk_dist_registry):
attributes = AccessAttributes(
roles=["admin", "researcher"],
teams=["ai-team", "ml-team"],
@ -137,9 +114,9 @@ async def test_registry_serialization(registry):
access_attributes=attributes,
)
await registry.register(model)
await cached_disk_dist_registry.register(model)
new_registry = CachedDiskDistributionRegistry(registry.kvstore)
new_registry = CachedDiskDistributionRegistry(cached_disk_dist_registry.kvstore)
await new_registry.initialize()
loaded_model = await new_registry.get("model", "model-serialize")