""" Tests for POST /generate — LiteLLM gateway. LiteLLM is mocked; no real network calls. """ import json import pytest import httpx from unittest.mock import AsyncMock, patch from httpx import AsyncClient, ASGITransport, Response from main import app, PromptContext from prompts import PROMPTS, get_prompt _build_user_v1 = PROMPTS["v1"].build_user def _litellm_response(candidates: list[dict]) -> Response: import httpx body = { "model": "tip-generator", "choices": [{"message": {"content": json.dumps(candidates)}}], "usage": {"prompt_tokens": 10, "completion_tokens": 20}, } req = httpx.Request("POST", "http://litellm/chat/completions") return Response(200, json=body, request=req) @pytest.mark.anyio async def test_generate_returns_candidates(): fake_items = [ {"id": "tip-1", "content": "Do the overdue task now.", "rationale": "It's been waiting."}, {"id": "tip-2", "content": "Take a 5-minute break.", "rationale": "You've been working long."}, ] mock_resp = _litellm_response(fake_items) with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(return_value=mock_resp) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post("/generate", json={"user_id": "u1", "n": 2}) assert resp.status_code == 200 data = resp.json() assert len(data["candidates"]) == 2 assert data["candidates"][0]["id"] == "tip-1" assert data["model"] == "tip-generator" @pytest.mark.anyio async def test_generate_strips_markdown_fence(): fake_items = [{"id": "tip-a", "content": "Focus.", "rationale": "Now."}] fenced = "```json\n" + json.dumps(fake_items) + "\n```" body = { "model": "tip-generator", "choices": [{"message": {"content": fenced}}], "usage": {}, } req = httpx.Request("POST", "http://litellm/chat/completions") mock_resp = Response(200, json=body, request=req) with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(return_value=mock_resp) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post("/generate", json={"user_id": "u1"}) assert resp.status_code == 200 assert resp.json()["candidates"][0]["id"] == "tip-a" @pytest.mark.anyio async def test_generate_503_on_unreachable(): import httpx as _httpx with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(side_effect=_httpx.ConnectError("refused")) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post("/generate", json={"user_id": "u1"}) assert resp.status_code == 503 def test_build_prompt_includes_tasks(): ctx = PromptContext( tasks=[{"content": "Write report", "priority": 4, "is_overdue": True, "due_date": "2026-04-15"}], hour_of_day=9, day_of_week=2, ) prompt = _build_user_v1(ctx, n=3) assert "Write report" in prompt assert "09:00" in prompt assert "Generate 3 tips" in prompt def test_build_prompt_truncates_at_five(): tasks = [{"content": f"Task {i}", "priority": 1, "is_overdue": False, "due_date": None} for i in range(8)] ctx = PromptContext(tasks=tasks, hour_of_day=12) prompt = _build_user_v1(ctx, n=2) assert "Task 4" in prompt assert "Task 5" not in prompt def test_build_prompt_extra_fields(): ctx = PromptContext(tasks=[], hour_of_day=8, extra={"mood": "focused", "energy": "high"}) prompt = _build_user_v1(ctx, n=1) assert "mood: focused" in prompt assert "energy: high" in prompt def test_build_prompt_empty_tasks_no_task_line(): ctx = PromptContext(tasks=[], hour_of_day=10) prompt = _build_user_v1(ctx, n=2) assert "Tasks:" not in prompt assert "Generate 2 tips" in prompt def test_build_prompt_tasks_sorted_overdue_first(): tasks = [ {"content": "Low priority", "priority": 1, "is_overdue": False, "task_age_days": 0}, {"content": "Overdue task", "priority": 2, "is_overdue": True, "task_age_days": 3}, ] ctx = PromptContext(tasks=tasks, hour_of_day=9) prompt = _build_user_v1(ctx, n=2) assert prompt.index("Overdue task") < prompt.index("Low priority") def test_build_prompt_includes_profile_features(): ctx = PromptContext( tasks=[], hour_of_day=14, profile_features={"completion_rate_30d": 0.75, "dismiss_rate_30d": 0.1, "preferred_hour": 9}, ) prompt = _build_user_v1(ctx, n=1) assert "User profile:" in prompt assert "completion_rate=75%" in prompt assert "dismiss_rate=10%" in prompt assert "preferred_hour=09:00" in prompt def test_build_prompt_no_profile_line_when_empty(): ctx = PromptContext(tasks=[], hour_of_day=10, profile_features={}) prompt = _build_user_v1(ctx, n=1) assert "User profile:" not in prompt def test_build_prompt_profile_partial_fields(): ctx = PromptContext( tasks=[], hour_of_day=10, profile_features={"completion_rate_30d": 0.5}, ) prompt = _build_user_v1(ctx, n=1) assert "completion_rate=50%" in prompt assert "dismiss_rate" not in prompt @pytest.mark.anyio async def test_generate_retry_succeeds_on_second_attempt(): """First response is invalid JSON; second is valid. Should return 200.""" valid_items = [{"id": "tip-ok", "content": "Retry worked.", "rationale": "Second try."}] bad_req = httpx.Request("POST", "http://litellm/chat/completions") bad_resp = Response(200, json={ "model": "tip-generator", "choices": [{"message": {"content": "this is not json"}}], "usage": {}, }, request=bad_req) good_resp = _litellm_response(valid_items) with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(side_effect=[bad_resp, good_resp]) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post("/generate", json={"user_id": "u1", "n": 1}) assert resp.status_code == 200 assert resp.json()["candidates"][0]["id"] == "tip-ok" assert instance.post.call_count == 2 # Retry message should include the correction suffix second_call_messages = instance.post.call_args_list[1][1]["json"]["messages"] assert any("not valid JSON" in m["content"] for m in second_call_messages) @pytest.mark.anyio async def test_generate_502_after_all_retries_exhausted(): """All attempts return invalid JSON → 502.""" bad_req = httpx.Request("POST", "http://litellm/chat/completions") def _bad_resp(): return Response(200, json={ "model": "tip-generator", "choices": [{"message": {"content": "not json at all"}}], "usage": {}, }, request=bad_req) from main import _MAX_GENERATE_RETRIES responses = [_bad_resp() for _ in range(1 + _MAX_GENERATE_RETRIES)] with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(side_effect=responses) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post("/generate", json={"user_id": "u1"}) assert resp.status_code == 502 assert "retries" in resp.json()["detail"] @pytest.mark.anyio async def test_generate_502_on_upstream_http_error(): """LiteLLM returns 500 → HTTPStatusError → 502.""" err_req = httpx.Request("POST", "http://litellm/chat/completions") err_resp = Response(500, text="internal error", request=err_req) with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(side_effect=httpx.HTTPStatusError( "500", request=err_req, response=err_resp )) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post("/generate", json={"user_id": "u1"}) assert resp.status_code == 502 assert "LiteLLM error" in resp.json()["detail"] def test_parse_llm_json_bare_fence(): from main import _parse_llm_json raw = "```\n[{\"id\":\"x\",\"content\":\"hi\"}]\n```" items = _parse_llm_json(raw) assert items[0]["id"] == "x" def test_parse_llm_json_no_fence(): from main import _parse_llm_json raw = '[{"id":"plain","content":"no fence"}]' items = _parse_llm_json(raw) assert items[0]["id"] == "plain" def test_parse_llm_json_raises_on_invalid(): from main import _parse_llm_json with pytest.raises((ValueError, Exception)): _parse_llm_json("this is not json") # ── Prompt registry / selection (#84) ────────────────────────────────────── def test_prompt_registry_contains_expected_versions(): assert set(PROMPTS.keys()) >= {"v1", "v2-mentor", "v3-few-shot"} # v2-mentor must differ from v1 in tone — easiest assertion: different system prompt. assert PROMPTS["v1"].system != PROMPTS["v2-mentor"].system # v3-few-shot must include curated example content in its system prompt. assert "Examples" in PROMPTS["v3-few-shot"].system def test_get_prompt_unknown_raises_keyerror(): with pytest.raises(KeyError): get_prompt("does-not-exist") def test_get_prompt_default_when_none(): p = get_prompt(None) assert p.version == "v1" # current DEFAULT_PROMPT_VERSION @pytest.mark.anyio async def test_generate_echoes_selected_prompt_version(): """Server should report back which prompt_version it actually used.""" fake_items = [{"id": "tip-1", "content": "x", "rationale": "y"}] mock_resp = _litellm_response(fake_items) with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(return_value=mock_resp) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post( "/generate", json={"user_id": "u1", "n": 1, "prompt_version": "v2-mentor"}, ) assert resp.status_code == 200 assert resp.json()["prompt_version"] == "v2-mentor" @pytest.mark.anyio async def test_generate_passes_profile_features_to_prompt(): """profile_features from GenerateRequest should appear in the user message sent to LiteLLM.""" fake_items = [{"id": "tip-1", "content": "x", "rationale": "y"}] mock_resp = _litellm_response(fake_items) captured_payload: list[dict] = [] async def _capture(url, *, json, headers): captured_payload.append(json) return mock_resp with patch("main.httpx.AsyncClient") as MockClient: instance = AsyncMock() instance.post = AsyncMock(side_effect=_capture) instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) MockClient.return_value = instance async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post("/generate", json={ "user_id": "u1", "n": 1, "profile_features": {"completion_rate_30d": 0.8, "preferred_hour": 10}, }) assert resp.status_code == 200 user_msg = captured_payload[0]["messages"][1]["content"] assert "User profile:" in user_msg assert "completion_rate=80%" in user_msg assert "preferred_hour=10:00" in user_msg @pytest.mark.anyio async def test_generate_422_on_unknown_prompt_version(): async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: resp = await client.post( "/generate", json={"user_id": "u1", "n": 1, "prompt_version": "nonsense"}, ) assert resp.status_code == 422 assert "Unknown prompt_version" in resp.json()["detail"]