refactor: extract StarletteUserProxy and add comprehensive route utility tests

Author: guglielmo-san

src/a2a/server/routes/common.py

@@ -1,6 +1,13 @@
 from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
-from starlette.requests import Request
+if TYPE_CHECKING:
+    from starlette.requests import Request
+else:
+    try:
+        from starlette.requests import Request
+    except ImportError:
+        Request = Any
 
 from a2a.auth.user import UnauthenticatedUser, User
 from a2a.extensions.common import (
@@ -13,20 +20,27 @@
 UserBuilder = Callable[[Request], User]
 
 
-def default_user_builder(request: Request) -> User:
-    """Default strategy for creating an A2AUser from a Starlette Request."""
-    if 'user' in request.scope:
+class StarletteUser(User):
+    """Adapts a Starlette BaseUser to the A2A User interface."""
+
+    def __init__(self, user: Any):
+        self._user = user
 
-        class BaseUser(User):
-            @property
-            def is_authenticated(self) -> bool:
-                return request.user.is_authenticated
+    @property
+    def is_authenticated(self) -> bool:
+        """Returns whether the current user is authenticated."""
+        return self._user.is_authenticated
 
-            @property
-            def user_name(self) -> str:
-                return request.user.display_name
+    @property
+    def user_name(self) -> str:
+        """Returns the user name of the current user."""
+        return self._user.display_name
 
-        return BaseUser()
+
+def default_user_builder(request: Request) -> User:
+    """Default strategy for creating an A2AUser from a Starlette Request."""
+    if 'user' in request.scope:
+        return StarletteUser(request.user)
     return UnauthenticatedUser()
 
 
@@ -37,7 +51,7 @@ def build_server_call_context(
 
     Args:
         request: The incoming Starlette Request object.
-        user_builder: Optional custom user builder.
+        user_builder: A callable that creates a User from the request.
 
     Returns:
         A ServerCallContext instance populated with user and state.

tests/server/routes/test_common.py

@@ -0,0 +1,156 @@
+from unittest.mock import MagicMock
+
+import pytest
+from starlette.datastructures import Headers
+
+try:
+    from starlette.authentication import BaseUser as StarletteBaseUser
+except ImportError:
+    StarletteBaseUser = MagicMock()  # type: ignore
+
+from a2a.auth.user import UnauthenticatedUser
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
+from a2a.server.context import ServerCallContext
+from a2a.server.routes.common import (
+    StarletteUser,
+    build_server_call_context,
+    default_user_builder,
+)
+
+
+# --- StarletteUser Tests ---
+
+
+class TestStarletteUser:
+    def test_is_authenticated_true(self):
+        starlette_user = MagicMock(spec=StarletteBaseUser)
+        starlette_user.is_authenticated = True
+        proxy = StarletteUser(starlette_user)
+        assert proxy.is_authenticated is True
+
+    def test_is_authenticated_false(self):
+        starlette_user = MagicMock(spec=StarletteBaseUser)
+        starlette_user.is_authenticated = False
+        proxy = StarletteUser(starlette_user)
+        assert proxy.is_authenticated is False
+
+    def test_user_name(self):
+        starlette_user = MagicMock(spec=StarletteBaseUser)
+        starlette_user.display_name = 'Test User'
+        proxy = StarletteUser(starlette_user)
+        assert proxy.user_name == 'Test User'
+
+    def test_user_name_raises_attribute_error(self):
+        starlette_user = MagicMock(spec=StarletteBaseUser)
+        del starlette_user.display_name
+        proxy = StarletteUser(starlette_user)
+        with pytest.raises(AttributeError, match='display_name'):
+            _ = proxy.user_name
+
+
+# --- default_user_builder Tests ---
+
+
+def _make_mock_request(scope=None, headers=None):
+    request = MagicMock()
+    request.scope = scope or {}
+    request.headers = Headers(headers or {})
+    return request
+
+
+class TestDefaultUserBuilder:
+    def test_returns_unauthenticated_user_when_no_user_in_scope(self):
+        request = _make_mock_request(scope={})
+        user = default_user_builder(request)
+        assert isinstance(user, UnauthenticatedUser)
+        assert user.is_authenticated is False
+        assert user.user_name == ''
+
+    def test_returns_proxy_when_user_in_scope(self):
+        starlette_user = MagicMock()
+        starlette_user.is_authenticated = True
+        starlette_user.display_name = 'Alice'
+        request = _make_mock_request(scope={'user': starlette_user})
+        request.user = starlette_user
+
+        user = default_user_builder(request)
+        assert isinstance(user, StarletteUser)
+        assert user.is_authenticated is True
+        assert user.user_name == 'Alice'
+
+    def test_returns_unauthenticated_proxy_when_user_not_authenticated(self):
+        starlette_user = MagicMock()
+        starlette_user.is_authenticated = False
+        starlette_user.display_name = ''
+        request = _make_mock_request(scope={'user': starlette_user})
+        request.user = starlette_user
+
+        user = default_user_builder(request)
+        assert isinstance(user, StarletteUser)
+        assert user.is_authenticated is False
+
+
+# --- build_server_call_context Tests ---
+
+
+class TestBuildServerCallContext:
+    def test_basic_context_with_default_user_builder(self):
+        request = _make_mock_request(
+            scope={}, headers={'content-type': 'application/json'}
+        )
+        ctx = build_server_call_context(request, default_user_builder)
+
+        assert isinstance(ctx, ServerCallContext)
+        assert isinstance(ctx.user, UnauthenticatedUser)
+        assert 'headers' in ctx.state
+        assert ctx.state['headers']['content-type'] == 'application/json'
+        assert 'auth' not in ctx.state
+
+    def test_auth_populated_when_in_scope(self):
+        auth_credentials = MagicMock()
+        request = _make_mock_request(scope={'auth': auth_credentials})
+        request.auth = auth_credentials
+
+        ctx = build_server_call_context(request, default_user_builder)
+        assert ctx.state['auth'] is auth_credentials
+
+    def test_auth_not_populated_when_not_in_scope(self):
+        request = _make_mock_request(scope={})
+        ctx = build_server_call_context(request, default_user_builder)
+        assert 'auth' not in ctx.state
+
+    def test_headers_captured_in_state(self):
+        request = _make_mock_request(
+            headers={'x-custom': 'value', 'authorization': 'Bearer tok'}
+        )
+        ctx = build_server_call_context(request, default_user_builder)
+        assert ctx.state['headers']['x-custom'] == 'value'
+        assert ctx.state['headers']['authorization'] == 'Bearer tok'
+
+    def test_requested_extensions_single(self):
+        request = _make_mock_request(headers={HTTP_EXTENSION_HEADER: 'foo'})
+        ctx = build_server_call_context(request, default_user_builder)
+        assert ctx.requested_extensions == {'foo'}
+
+    def test_requested_extensions_comma_separated(self):
+        request = _make_mock_request(
+            headers={HTTP_EXTENSION_HEADER: 'foo, bar'}
+        )
+        ctx = build_server_call_context(request, default_user_builder)
+        assert ctx.requested_extensions == {'foo', 'bar'}
+
+    def test_no_extensions(self):
+        request = _make_mock_request()
+        ctx = build_server_call_context(request, default_user_builder)
+        assert ctx.requested_extensions == set()
+
+    def test_custom_user_builder(self):
+        custom_user = MagicMock(spec=UnauthenticatedUser)
+        custom_user.is_authenticated = True
+
+        def my_builder(req):
+            return custom_user
+
+        request = _make_mock_request()
+        ctx = build_server_call_context(request, my_builder)
+        assert ctx.user is custom_user