Mercurial

annotate mrjunejune/inference/litellm_gateway_test.py @ 261:b401627fc49e

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Add JRPG mock flows and interactive previews Add scripted mock SSE commands, custom event forwarding, animated chat turns, full-height message navigation, and a cyberpunk resume dossier. Co-authored-by: Copilot <[email protected]>
author MrJuneJune <mrjunejune@users.noreply.github.com>
date Wed, 05 Aug 2026 20:38:32 -0700
parents 1f9877b637e9
children
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
260
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
1 import os
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
2 import pathlib
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
3 import unittest
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
4 from unittest import mock
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
5
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
6 from fastapi.testclient import TestClient
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
7
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
8 from mrjunejune.inference.litellm_proxy import (
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
9 create_app,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
10 resolve_token_directory,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
11 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
12
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
13
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
14 class Chunk:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
15 def __init__(self, content):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
16 self.content = content
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
17
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
18 def model_dump(self, exclude_none=True):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
19 del exclude_none
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
20 return {
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
21 "choices": [
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
22 {"index": 0, "delta": {"content": self.content}}
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
23 ]
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
24 }
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
25
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
26
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
27 class Stream:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
28 def __init__(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
29 self.chunks = iter([Chunk("hello "), Chunk("traveler")])
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
30
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
31 def __aiter__(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
32 return self
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
33
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
34 async def __anext__(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
35 try:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
36 return next(self.chunks)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
37 except StopIteration as error:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
38 raise StopAsyncIteration from error
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
39
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
40
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
41 class LiteLlmGatewayTest(unittest.TestCase):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
42 def setUp(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
43 os.environ["LITELLM_MASTER_KEY"] = "test-key"
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
44 self.calls = []
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
45
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
46 async def completion(**payload):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
47 self.calls.append(payload)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
48 return Stream()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
49
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
50 self.client = TestClient(create_app(completion))
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
51
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
52 def test_health_does_not_require_provider_auth(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
53 response = self.client.get("/health/liveliness")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
54 self.assertEqual(response.status_code, 200)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
55 self.assertEqual(response.json(), {"status": "ready"})
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
56
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
57 def test_stream_requires_key_and_maps_alias(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
58 denied = self.client.post(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
59 "/v1/chat/completions",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
60 json={"model": "jrpg-copilot", "messages": [], "stream": True},
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
61 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
62 self.assertEqual(denied.status_code, 401)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
63
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
64 with self.client.stream(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
65 "POST",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
66 "/v1/chat/completions",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
67 headers={"Authorization": "Bearer test-key"},
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
68 json={
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
69 "model": "jrpg-copilot",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
70 "messages": [{"role": "user", "content": "hello"}],
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
71 "stream": True,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
72 "max_tokens": 99999,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
73 },
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
74 ) as response:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
75 body = "".join(response.iter_text())
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
76 self.assertEqual(response.status_code, 200)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
77 self.assertIn("hello ", body)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
78 self.assertIn("traveler", body)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
79 self.assertIn("data: [DONE]", body)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
80 self.assertEqual(self.calls[0]["model"], "github_copilot/gpt-4")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
81 self.assertEqual(self.calls[0]["max_tokens"], 1024)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
82
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
83 def test_token_directory_precedence(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
84 with mock.patch.dict(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
85 os.environ,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
86 {
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
87 "GITHUB_COPILOT_TOKEN_DIR": "/tmp/copilot-env",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
88 "MRJUNEJUNE_INFERENCE_STATE": "/tmp/inference-state",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
89 },
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
90 clear=True,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
91 ):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
92 self.assertEqual(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
93 resolve_token_directory("/tmp/copilot-argument"),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
94 pathlib.Path("/tmp/copilot-argument"),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
95 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
96 self.assertEqual(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
97 resolve_token_directory(),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
98 pathlib.Path("/tmp/copilot-env"),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
99 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
100
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
101 def test_token_directory_uses_persistent_state_root(self):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
102 with mock.patch.dict(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
103 os.environ,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
104 {"MRJUNEJUNE_INFERENCE_STATE": "/tmp/inference-state"},
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
105 clear=True,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
106 ):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
107 self.assertEqual(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
108 resolve_token_directory(),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
109 pathlib.Path("/tmp/inference-state/litellm-copilot"),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
110 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
111
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
112
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
113 if __name__ == "__main__":
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
114 unittest.main()