diff --git a/llama_stack/distribution/access_control/access_control.py b/llama_stack/distribution/access_control/access_control.py index 89ae60343..075152ce4 100644 --- a/llama_stack/distribution/access_control/access_control.py +++ b/llama_stack/distribution/access_control/access_control.py @@ -105,7 +105,17 @@ def is_action_allowed( return False -def build_access_denied_message(action: str | None, resource: ProtectedResource | None, user: User | None) -> str: +class AccessDeniedError(RuntimeError): + def __init__(self, action: str | None = None, resource: ProtectedResource | None = None, user: User | None = None): + self.action = action + self.resource = resource + self.user = user + + message = _build_access_denied_message(action, resource, user) + super().__init__(message) + + +def _build_access_denied_message(action: str | None, resource: ProtectedResource | None, user: User | None) -> str: """Build detailed error message for access denied scenarios.""" if action and resource and user: resource_info = f"{resource.type}::{resource.identifier}" @@ -119,13 +129,3 @@ def build_access_denied_message(action: str | None, resource: ProtectedResource message = "Insufficient permissions" return message - - -class AccessDeniedError(RuntimeError): - def __init__(self, action: str | None = None, resource: ProtectedResource | None = None, user: User | None = None): - self.action = action - self.resource = resource - self.user = user - - message = build_access_denied_message(action, resource, user) - super().__init__(message) diff --git a/tests/unit/fixtures.py b/tests/unit/fixtures.py index 7174d2e78..4e50c5e08 100644 --- a/tests/unit/fixtures.py +++ b/tests/unit/fixtures.py @@ -4,14 +4,14 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import pytest +import pytest_asyncio from llama_stack.distribution.store.registry import CachedDiskDistributionRegistry, DiskDistributionRegistry from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl -@pytest.fixture(scope="function") +@pytest_asyncio.fixture(scope="function") async def sqlite_kvstore(tmp_path): db_path = tmp_path / "test_kv.db" kvstore_config = SqliteKVStoreConfig(db_path=db_path.as_posix()) @@ -20,14 +20,14 @@ async def sqlite_kvstore(tmp_path): yield kvstore -@pytest.fixture(scope="function") +@pytest_asyncio.fixture(scope="function") async def disk_dist_registry(sqlite_kvstore): registry = DiskDistributionRegistry(sqlite_kvstore) await registry.initialize() yield registry -@pytest.fixture(scope="function") +@pytest_asyncio.fixture(scope="function") async def cached_disk_dist_registry(sqlite_kvstore): registry = CachedDiskDistributionRegistry(sqlite_kvstore) await registry.initialize() diff --git a/tests/unit/providers/agents/test_persistence_access_control.py b/tests/unit/providers/agents/test_persistence_access_control.py index d5b876a09..656d1e53c 100644 --- a/tests/unit/providers/agents/test_persistence_access_control.py +++ b/tests/unit/providers/agents/test_persistence_access_control.py @@ -9,6 +9,7 @@ from datetime import datetime from unittest.mock import patch import pytest +import pytest_asyncio from llama_stack.apis.agents import Turn from llama_stack.apis.inference import CompletionMessage, StopReason @@ -16,7 +17,7 @@ from llama_stack.distribution.datatypes import User from llama_stack.providers.inline.agents.meta_reference.persistence import AgentPersistence, AgentSessionInfo -@pytest.fixture +@pytest_asyncio.fixture async def test_setup(sqlite_kvstore): agent_persistence = AgentPersistence(agent_id="test_agent", kvstore=sqlite_kvstore, policy={}) yield agent_persistence diff --git a/tests/unit/providers/vector_io/test_sqlite_vec.py b/tests/unit/providers/vector_io/test_sqlite_vec.py index bbac717c7..5d9d92cf3 100644 --- a/tests/unit/providers/vector_io/test_sqlite_vec.py +++ b/tests/unit/providers/vector_io/test_sqlite_vec.py @@ -148,7 +148,7 @@ async def test_chunk_id_conflict(sqlite_vec_index, sample_chunks, embedding_dime assert len(chunk_ids) == len(set(chunk_ids)), "Duplicate chunk IDs detected across batches!" -@pytest.fixture(scope="session") +@pytest_asyncio.fixture(scope="session") async def sqlite_vec_adapter(sqlite_connection): config = type("Config", (object,), {"db_path": ":memory:"}) # Mock config with in-memory database adapter = SQLiteVecVectorIOAdapter(config=config, inference_api=None) diff --git a/tests/unit/server/test_access_control.py b/tests/unit/server/test_access_control.py index f9ad47b0c..af03ddacb 100644 --- a/tests/unit/server/test_access_control.py +++ b/tests/unit/server/test_access_control.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, Mock, patch import pytest +import pytest_asyncio import yaml from pydantic import TypeAdapter, ValidationError @@ -26,7 +27,7 @@ def _return_model(model): return model -@pytest.fixture +@pytest_asyncio.fixture async def test_setup(cached_disk_dist_registry): mock_inference = Mock() mock_inference.__provider_spec__ = MagicMock() @@ -245,7 +246,7 @@ async def test_automatic_access_attributes(mock_get_authenticated_user, test_set assert model.identifier == "auto-access-model" -@pytest.fixture +@pytest_asyncio.fixture async def test_setup_with_access_policy(cached_disk_dist_registry): mock_inference = Mock() mock_inference.__provider_spec__ = MagicMock()