aws-samples / aws-samples/sample-agentic-platform

bug: `tests/unit/core/client/test_litellm_gateway_client.py`: 4 tests assert an outdated `LiteLLMGatewayClient` API

Open
#88 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.