aws-samples / aws-samples/sample-agentic-platform
bug: `tests/unit/core/client/test_litellm_gateway_client.py`: 4 tests assert an outdated `LiteLLMGatewayClient` API
- Dominant language
- Python
- Stars
- 133
- Forks
- 54
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
Four tests in `test_litellm_gateway_client.py` assert behaviour that `LiteLLMGatewayClient` no longer implements. These appear to be stale tests trailing a refactor rather than defects in the client.
Note: these are only reachable after the collection errors in `tests/unit/agent/{diy_agent,pydanticai_agent}/` are resolved (filed separately).
## Version
- Commit: `e41eca1` (`main`, 2026-06-17)
- Python 3.12.8, pytest 8.3.5, macOS arm64
- No local modifications
## Reproduction
```bash
uv run pytest tests/unit/core/client/test_litellm_gateway_client.py -q
```
## Actual result
`4 failed`:
| Test | Asserts | Current behaviour |
| --- | --- | --- |
| `test_init_with_env_vars` | `api_key == 'env-key'` | `api_key is None` |
| `test_get_headers_with_auth_token` | `Authorization: Bearer context-token` | `Bearer ` |
| `test_get_client` | returns a `dict` | returns `LiteLLMClientInfo` |
| `test_get_openai_client` | method exists | `AttributeError` |
## Analysis
Each mismatch traces to a deliberate change in the client.
**`test_init_with_env_vars`** — `LITELLM_API_KEY` is read at module import time:
```python
LITELLM_API_KEY = os.getenv('LITELLM_KEY')
```
Patching `os.environ` inside a test therefore has no effect on `__init__`. The test would need to patch the module attribute, or the client would need to read the environment lazily.
**`test_get_client`** — now returns a `LiteLLMClientInfo` pydantic model (`api_key`, `api_endpoint`) instead of a plain dict.
**`test_get_openai_client`** — no such method exists on the client.
**`test_get_headers_with_auth_token`** — the test expects the caller's JWT from `request_context` to be forwarded. `_get_headers()` uses the configured `self.api_key` instead. This looks correct as written: unlike `memory_gateway_client` and `retrieval_gateway_client`, which forward the caller's token to internal services that validate it, LiteLLM authenticates with its own key. `infrastructure/modules/litellm/agent-secret.tf` provisions that key as `LITELLM_KEY` specifically for agents, and `litellm_config.yaml` reads `master_key: os.environ/LITELLM_MASTER_KEY`. Forwarding a Cognito JWT to LiteLLM would not authenticate. The test appears to encode earlier behaviour.
## Suggested fix
Update the four tests to match the current API. A maintainer should confirm the intended auth model for the LLM gateway before changing `test_get_headers_with_auth_token`, since that test is the only remaining assertion that per-request tokens reach LiteLLM.
Minor cleanup while in the file: `litellm_gateway_client.py` imports `get_auth_token` and carries the comment "Try to get auth token from context, fall back to configured API key", but never calls it. If the static-key model is correct, the import and comment are dead and actively misleading:
```python
from agentic_platform.core.context.request_context import get_auth_token
...
def _get_headers(self) -> Dict[str, str]:
# Try to get auth token from context, fall back to configured API key
auth_token = self.api_key # get_auth_token() never called
```
Contributor guide
Research direction
Run `uv run pytest tests/unit/core/client/test_litellm_gateway_client.py -q`, noting that collection errors in `tests/unit/agent/{diy_agent,pydanticai_agent}/` currently block these tests. Read `test_litellm_gateway_client.py` with the current `LiteLLMGatewayClient`, `infrastructure/modules/litellm/agent-secret.tf`, and `litellm_config.yaml`; done means the four assertions match the confirmed API and auth model.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend-api-design, testing
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100