fix(llm-gateway): improve type safety and datetime consistency
- Add type annotations for mypy compliance - Use UTC-aware datetimes consistently (datetime.now(UTC)) - Add type: ignore comments for LiteLLM incomplete stubs - Fix import ordering and formatting - Update pyproject.toml mypy configuration 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -108,19 +108,19 @@ class TestCostTracker:
|
||||
)
|
||||
|
||||
# Verify by getting usage report
|
||||
report = asyncio.run(
|
||||
tracker.get_project_usage("proj-123", period="day")
|
||||
)
|
||||
report = asyncio.run(tracker.get_project_usage("proj-123", period="day"))
|
||||
|
||||
assert report.total_requests == 1
|
||||
assert report.total_cost_usd == pytest.approx(0.01, rel=0.01)
|
||||
|
||||
def test_record_usage_disabled(self, tracker_settings: Settings) -> None:
|
||||
"""Test recording is skipped when disabled."""
|
||||
settings = Settings(**{
|
||||
**tracker_settings.model_dump(),
|
||||
"cost_tracking_enabled": False,
|
||||
})
|
||||
settings = Settings(
|
||||
**{
|
||||
**tracker_settings.model_dump(),
|
||||
"cost_tracking_enabled": False,
|
||||
}
|
||||
)
|
||||
fake_redis = fakeredis.aioredis.FakeRedis(decode_responses=True)
|
||||
disabled_tracker = CostTracker(redis_client=fake_redis, settings=settings)
|
||||
|
||||
@@ -157,9 +157,7 @@ class TestCostTracker:
|
||||
)
|
||||
|
||||
# Verify session usage
|
||||
session_usage = asyncio.run(
|
||||
tracker.get_session_usage("session-789")
|
||||
)
|
||||
session_usage = asyncio.run(tracker.get_session_usage("session-789"))
|
||||
|
||||
assert session_usage["session_id"] == "session-789"
|
||||
assert session_usage["total_cost_usd"] == pytest.approx(0.01, rel=0.01)
|
||||
@@ -199,9 +197,7 @@ class TestCostTracker:
|
||||
)
|
||||
)
|
||||
|
||||
report = asyncio.run(
|
||||
tracker.get_project_usage("proj-123", period="day")
|
||||
)
|
||||
report = asyncio.run(tracker.get_project_usage("proj-123", period="day"))
|
||||
|
||||
assert report.total_requests == 2
|
||||
assert len(report.by_model) == 2
|
||||
@@ -221,9 +217,7 @@ class TestCostTracker:
|
||||
)
|
||||
)
|
||||
|
||||
report = asyncio.run(
|
||||
tracker.get_agent_usage("agent-456", period="day")
|
||||
)
|
||||
report = asyncio.run(tracker.get_agent_usage("agent-456", period="day"))
|
||||
|
||||
assert report.entity_id == "agent-456"
|
||||
assert report.entity_type == "agent"
|
||||
@@ -243,12 +237,8 @@ class TestCostTracker:
|
||||
)
|
||||
|
||||
# Check different periods
|
||||
hour_report = asyncio.run(
|
||||
tracker.get_project_usage("proj-123", period="hour")
|
||||
)
|
||||
day_report = asyncio.run(
|
||||
tracker.get_project_usage("proj-123", period="day")
|
||||
)
|
||||
hour_report = asyncio.run(tracker.get_project_usage("proj-123", period="hour"))
|
||||
day_report = asyncio.run(tracker.get_project_usage("proj-123", period="day"))
|
||||
month_report = asyncio.run(
|
||||
tracker.get_project_usage("proj-123", period="month")
|
||||
)
|
||||
@@ -306,9 +296,7 @@ class TestCostTracker:
|
||||
|
||||
def test_check_budget_default_limit(self, tracker: CostTracker) -> None:
|
||||
"""Test budget check with default limit."""
|
||||
within, current, limit = asyncio.run(
|
||||
tracker.check_budget("proj-123")
|
||||
)
|
||||
within, current, limit = asyncio.run(tracker.check_budget("proj-123"))
|
||||
|
||||
assert limit == 1000.0 # Default from settings
|
||||
|
||||
|
||||
Reference in New Issue
Block a user