provider.rs

crates/cortex-core/tests/provider.rs

1. //! The provider client against the stand-in server: every path a real stream
2. //! can take, including cancelling mid-reply.
4. use std::time::Duration;
6. use cortex_core::provider::{ChatRequest, Client, OutMessage, ProviderError, StreamItem, StreamOptions};
7. use cortex_core::types::Usage;
8. use tokio_util::sync::CancellationToken;
10. fn req(model: &str, text: &str) -> ChatRequest {
11. ChatRequest {
12. model: model.into(),
13. messages: vec![OutMessage::text("system", "be brief"), OutMessage::text("user", text)],
14. stream: true,
15. stream_options: Some(StreamOptions { include_usage: true }),
16. temperature: None,
17. max_tokens: None,
18. }
19. }
21. async fn run(client: &Client, r: &ChatRequest) -> (cortex_core::provider::StreamOutcome, Vec<StreamItem>) {
22. let mut items = Vec::new();
23. let out = client.stream_chat(r, &CancellationToken::new(), |i| items.push(i)).await;
24. (out, items)
25. }
27. #[tokio::test]
28. async fn streams_text_and_final_usage() {
29. let s = standin::spawn().await;
30. let c = Client::new(&s.base_url(), Some("good-key".into()));
31. let (out, items) = run(&c, &req("standin/echo", "hello there friend")).await;
32. assert_eq!(out.error, None);
33. assert_eq!(out.content, "You said: hello there friend");
34. assert_eq!(out.finish_reason.as_deref(), Some("stop"));
35. // 2 + 3 words, plus 4 per message
36. assert_eq!(out.usage, Some(Usage { prompt_tokens: 13, completion_tokens: 5, cached_tokens: 0 }));
37. let deltas = items.iter().filter(|i| matches!(i, StreamItem::Content(_))).count();
38. assert_eq!(deltas, 5);
39. assert!(out.ttft_ms.is_some());
40. let rec = &s.requests()[0];
41. assert_eq!(rec.authorization.as_deref(), Some("Bearer good-key"));
42. assert_eq!(rec.body["stream_options"]["include_usage"], true);
43. assert_eq!(rec.body["stream"], true);
44. }
46. #[tokio::test]
47. async fn no_key_sends_no_auth_header() {
48. let s = standin::spawn().await;
49. let c = Client::new(&s.base_url(), Some(String::new()));
50. run(&c, &req("standin/echo", "hi")).await;
51. assert_eq!(s.requests()[0].authorization, None);
52. }
54. #[tokio::test]
55. async fn usage_variants() {
56. let s = standin::spawn().await;
57. let c = Client::new(&s.base_url(), None);
58. let (out, _) = run(&c, &req("standin/inline-usage", "a b")).await;
59. assert_eq!(out.usage.unwrap().completion_tokens, 4);
60. let (out, _) = run(&c, &req("standin/no-usage", "a b")).await;
61. assert_eq!(out.usage, None);
62. assert_eq!(out.content, "You said: a b");
63. let (out, _) = run(&c, &req("standin/cached", "one two three")).await;
64. assert_eq!(out.usage.unwrap().cached_tokens, 13);
65. assert_eq!(out.provider_cost_usd, Some(0.0000123));
66. }
68. #[tokio::test]
69. async fn reasoning_is_separate_from_content() {
70. let s = standin::spawn().await;
71. let c = Client::new(&s.base_url(), None);
72. let (out, items) = run(&c, &req("standin/reasoning", "why")).await;
73. assert_eq!(out.reasoning, "Checking the question.");
74. assert_eq!(out.content, "You said: why");
75. assert!(matches!(items[0], StreamItem::Reasoning(_)));
76. }
78. #[tokio::test]
79. async fn images_go_as_content_parts() {
80. let s = standin::spawn().await;
81. let c = Client::new(&s.base_url(), None);
82. let mut r = req("standin/vision", "");
83. r.messages[1] = OutMessage::with_images(
84. "user",
85. "what is in these",
86. vec!["data:image/png;base64,AAAA".into(), "data:image/jpeg;base64,BBBB".into()],
87. );
88. let (out, _) = run(&c, &r).await;
89. assert_eq!(out.content, "I see 2 images.");
90. let body = &s.requests()[0].body;
91. assert_eq!(body["messages"][1]["content"][2]["type"], "image_url");
92. }
94. #[tokio::test]
95. async fn http_errors_map_to_plain_errors() {
96. let s = standin::spawn().await;
97. let bad = Client::new(&s.base_url(), Some("bad-key".into()));
98. assert_eq!(run(&bad, &req("standin/echo", "x")).await.0.error, Some(ProviderError::InvalidKey));
99. // Saving a key checks it this way before storing it.
100. assert_eq!(bad.list_models().await.unwrap_err(), ProviderError::InvalidKey);
101. let c = Client::new(&s.base_url(), None);
102. assert_eq!(run(&c, &req("standin/402", "x")).await.0.error, Some(ProviderError::InsufficientBalance));
103. assert_eq!(
104. run(&c, &req("standin/429", "x")).await.0.error,
105. Some(ProviderError::RateLimited { retry_after: Some(7) })
106. );
107. assert_eq!(
108. run(&c, &req("standin/500", "x")).await.0.error,
109. Some(ProviderError::Upstream { status: 500, message: "upstream exploded".into() })
110. );
111. }
113. #[tokio::test]
114. async fn midstream_error_keeps_the_partial_reply() {
115. let s = standin::spawn().await;
116. let c = Client::new(&s.base_url(), None);
117. let (out, _) = run(&c, &req("standin/midstream-error", "one two three four")).await;
118. assert_eq!(out.content, "You said: one"); // three words, then the error
119. assert_eq!(out.error, Some(ProviderError::Stream("model overloaded".into())));
120. assert_eq!(out.usage, None);
121. }
123. #[tokio::test]
124. async fn unreachable_server() {
125. // Bind then drop a listener to get a port with nothing on it.
126. let port = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port();
127. let c = Client::new(&format!("http://127.0.0.1:{port}/v1"), None);
128. let (out, _) = run(&c, &req("standin/echo", "x")).await;
129. assert_eq!(out.error.unwrap().kind(), "unreachable");
130. }
132. #[tokio::test]
133. async fn cancel_drops_the_connection_and_keeps_what_arrived() {
134. let s = standin::spawn().await;
135. let c = Client::new(&s.base_url(), None);
136. let cancel = CancellationToken::new();
137. let mut deltas = 0;
138. let trigger = cancel.clone();
139. let out = c
140. .stream_chat(&req("standin/slow", "go"), &cancel, |i| {
141. if matches!(i, StreamItem::Content(_)) {
142. deltas += 1;
143. if deltas == 3 {
144. trigger.cancel();
145. }
146. }
147. })
148. .await;
149. assert!(out.cancelled);
150. assert_eq!(out.error, None);
151. assert!(out.content.starts_with("word1 word2"), "{}", out.content);
152. assert!(out.content.len() < 60);
153. // The server notices the dropped connection on its next write.
154. let mut seen = 0;
155. for _ in 0..40 {
156. seen = s.disconnects();
157. if seen > 0 {
158. break;
159. }
160. tokio::time::sleep(Duration::from_millis(50)).await;
161. }
162. assert_eq!(seen, 1, "server never saw the client go away");
163. }
165. #[tokio::test]
166. async fn cancel_before_the_response_arrives() {
167. let s = standin::spawn().await;
168. let c = Client::new(&s.base_url(), None);
169. let cancel = CancellationToken::new();
170. cancel.cancel();
171. let out = c.stream_chat(&req("standin/echo", "x"), &cancel, |_| {}).await;
172. assert!(out.cancelled && out.content.is_empty());
173. }
175. #[tokio::test]
176. async fn lists_models() {
177. let s = standin::spawn().await;
178. let c = Client::new(&s.base_url(), None);
179. let v = c.list_models().await.unwrap();
180. assert_eq!(v["data"][0]["id"], "standin/echo");
181. }