mirror of
https://github.com/hanzoai/insights.git
synced 2026-10-30 16:15:31 +00:00
The tree was a debranded snapshot of upstream taken around 2026-02-20, and upstream has moved about five months since. This replaces it with upstream master, mechanically re-derived: bin/debrand does the rename, and the pipeline and squashed migrations are restored from the shipped tree. Not a merge. The previous import replaced the tree wholesale in a single commit, so there is no shared content history to merge against. Reproducing the rename on the new tree is what actually converges. What survives: .hanzo/workflows, the Dockerfiles and bin/docker* entrypoints, compose files, LLM.md, and insights/migrations. The squashed initial migration is load-bearing — the live database records it by name, so upstream's ~918 individual migrations cannot come back. Nine products are dropped as unreferenced. Twenty-one more that we do not run are kept anyway because the core app still imports them; unpicking those is a refactor and does not belong in a rebase. This base does not build yet, and the reasons are not in the rename: upstream now wants Python 3.13.13 against our 3.12.12 and Django 5.2 against our 4.2, and the dependency manifests still name upstream packages that we republish under our own names. Those are decisions, not renames.
172 lines
5.4 KiB
Python
172 lines
5.4 KiB
Python
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from llm_gateway.rate_limiting.throttles import ThrottleContext
|
|
from llm_gateway.request_context import (
|
|
RequestContext,
|
|
apply_insights_context_from_headers,
|
|
get_insights_flags,
|
|
get_insights_properties,
|
|
record_cost,
|
|
request_context_var,
|
|
set_request_context,
|
|
set_throttle_context,
|
|
throttle_context_var,
|
|
throttle_runner_var,
|
|
)
|
|
|
|
|
|
def make_mock_user() -> MagicMock:
|
|
user = MagicMock()
|
|
user.user_id = 1
|
|
user.team_id = 1
|
|
user.application_id = None
|
|
user.auth_method = "api_key"
|
|
return user
|
|
|
|
|
|
class TestRecordCost:
|
|
@pytest.fixture(autouse=True)
|
|
def reset_context_vars(self) -> None:
|
|
throttle_runner_var.set(None)
|
|
throttle_context_var.set(None)
|
|
|
|
async def test_record_cost_calls_runner(self) -> None:
|
|
mock_runner = MagicMock()
|
|
mock_runner.record_cost = AsyncMock()
|
|
|
|
context = ThrottleContext(
|
|
user=make_mock_user(),
|
|
product="llm_gateway",
|
|
)
|
|
|
|
set_throttle_context(mock_runner, context)
|
|
|
|
await record_cost(0.0015)
|
|
|
|
mock_runner.record_cost.assert_called_once_with(context, 0.0015)
|
|
|
|
async def test_record_cost_no_op_without_context(self) -> None:
|
|
await record_cost(0.0015)
|
|
|
|
async def test_record_cost_no_op_with_partial_context(self) -> None:
|
|
mock_runner = MagicMock()
|
|
throttle_runner_var.set(mock_runner)
|
|
|
|
await record_cost(0.0015)
|
|
|
|
mock_runner.record_cost.assert_not_called()
|
|
|
|
async def test_record_cost_does_not_override_existing_end_user_id(self) -> None:
|
|
mock_runner = MagicMock()
|
|
mock_runner.record_cost = AsyncMock()
|
|
|
|
context = ThrottleContext(
|
|
user=make_mock_user(),
|
|
product="llm_gateway",
|
|
end_user_id="authenticated-user-42",
|
|
)
|
|
|
|
set_throttle_context(mock_runner, context)
|
|
|
|
await record_cost(0.0015, end_user_id="attacker-supplied-id")
|
|
|
|
assert context.end_user_id == "authenticated-user-42"
|
|
mock_runner.record_cost.assert_called_once_with(context, 0.0015)
|
|
|
|
async def test_record_cost_sets_end_user_id_when_not_already_set(self) -> None:
|
|
mock_runner = MagicMock()
|
|
mock_runner.record_cost = AsyncMock()
|
|
|
|
context = ThrottleContext(
|
|
user=make_mock_user(),
|
|
product="llm_gateway",
|
|
end_user_id=None,
|
|
)
|
|
|
|
set_throttle_context(mock_runner, context)
|
|
|
|
await record_cost(0.0015, end_user_id="fallback-id")
|
|
|
|
assert context.end_user_id == "fallback-id"
|
|
mock_runner.record_cost.assert_called_once_with(context, 0.0015)
|
|
|
|
|
|
class TestApplyInsightsContextFromHeaders:
|
|
@pytest.fixture(autouse=True)
|
|
def reset_request_context(self) -> None:
|
|
request_context_var.set(None)
|
|
|
|
def test_sets_properties_and_flags_from_headers(self) -> None:
|
|
request = MagicMock()
|
|
request.headers = MagicMock()
|
|
request.headers.items.return_value = [
|
|
("X-INSIGHTS-PROPERTY-VARIANT", "memes"),
|
|
("X-INSIGHTS-FLAG-EXPERIMENT", "test"),
|
|
("Content-Type", "application/json"),
|
|
]
|
|
set_request_context(RequestContext(request_id="req-123"))
|
|
|
|
apply_insights_context_from_headers(request)
|
|
|
|
assert get_insights_properties() == {"variant": "memes"}
|
|
assert get_insights_flags() == {"experiment": "test"}
|
|
|
|
def test_leaves_existing_context_unchanged_without_matching_headers(self) -> None:
|
|
request = MagicMock()
|
|
request.headers = MagicMock()
|
|
request.headers.items.return_value = [("Content-Type", "application/json")]
|
|
set_request_context(
|
|
RequestContext(
|
|
request_id="req-123",
|
|
insights_properties={"existing": "property"},
|
|
insights_flags={"existing": "flag"},
|
|
)
|
|
)
|
|
|
|
apply_insights_context_from_headers(request)
|
|
|
|
assert get_insights_properties() == {"existing": "property"}
|
|
assert get_insights_flags() == {"existing": "flag"}
|
|
|
|
def test_updates_only_properties_preserves_existing_flags(self) -> None:
|
|
request = MagicMock()
|
|
request.headers = MagicMock()
|
|
request.headers.items.return_value = [
|
|
("X-INSIGHTS-PROPERTY-VARIANT", "memes"),
|
|
("Content-Type", "application/json"),
|
|
]
|
|
set_request_context(
|
|
RequestContext(
|
|
request_id="req-123",
|
|
insights_properties={"existing": "property"},
|
|
insights_flags={"existing": "flag"},
|
|
)
|
|
)
|
|
|
|
apply_insights_context_from_headers(request)
|
|
|
|
assert get_insights_properties() == {"variant": "memes"}
|
|
assert get_insights_flags() == {"existing": "flag"}
|
|
|
|
def test_updates_only_flags_preserves_existing_properties(self) -> None:
|
|
request = MagicMock()
|
|
request.headers = MagicMock()
|
|
request.headers.items.return_value = [
|
|
("X-INSIGHTS-FLAG-EXPERIMENT", "test"),
|
|
("Content-Type", "application/json"),
|
|
]
|
|
set_request_context(
|
|
RequestContext(
|
|
request_id="req-123",
|
|
insights_properties={"existing": "property"},
|
|
insights_flags={"existing": "flag"},
|
|
)
|
|
)
|
|
|
|
apply_insights_context_from_headers(request)
|
|
|
|
assert get_insights_properties() == {"existing": "property"}
|
|
assert get_insights_flags() == {"experiment": "test"}
|