diff --git a/tests/unit/server/test_sse.py b/tests/unit/server/test_sse.py index c78122294..60e9f4609 100644 --- a/tests/unit/server/test_sse.py +++ b/tests/unit/server/test_sse.py @@ -5,10 +5,12 @@ # the root directory of this source tree. import asyncio +from unittest.mock import AsyncMock, MagicMock import pytest -from llama_stack.distribution.server.server import create_sse_event, sse_generator +from llama_stack.apis.common.responses import PaginatedResponse +from llama_stack.distribution.server.server import create_dynamic_typed_route, create_sse_event, sse_generator @pytest.mark.asyncio @@ -89,3 +91,24 @@ async def test_sse_generator_error_before_response_starts(): # We should have 1 error event assert len(seen_events) == 1 assert 'data: {"error":' in seen_events[0] + + +@pytest.mark.asyncio +async def test_paginated_response_url_setting(): + """Test that PaginatedResponse gets url set to route path.""" + + async def mock_api_method(): + return PaginatedResponse(data=[], has_more=False, url=None) + + route_handler = create_dynamic_typed_route(mock_api_method, "get", "/test/route") + + # Mock minimal request + request = MagicMock() + request.scope = {"user_attributes": {}, "principal": ""} + request.headers = {} + request.body = AsyncMock(return_value=b"") + + result = await route_handler(request) + + assert isinstance(result, PaginatedResponse) + assert result.url == "/test/route"