mockforge_http/
token_response.rs

1//! Token-based response resolution for HTTP handlers
2//!
3//! This module integrates the token resolver with HTTP response generation.
4
5use axum::{
6    body::Body,
7    http::{Response, StatusCode},
8    response::IntoResponse,
9};
10use mockforge_data::rag::RagConfig;
11use mockforge_data::{resolve_tokens, resolve_tokens_with_rag};
12use serde_json::Value;
13use tracing::*;
14
15/// Resolve tokens in a JSON response body
16pub async fn resolve_response_tokens(body: Value) -> Result<Value, String> {
17    resolve_tokens(&body)
18        .await
19        .map_err(|e| format!("Failed to resolve tokens: {}", e))
20}
21
22/// Resolve tokens in a JSON response body with RAG support
23pub async fn resolve_response_tokens_with_rag(
24    body: Value,
25    rag_config: RagConfig,
26) -> Result<Value, String> {
27    resolve_tokens_with_rag(&body, rag_config)
28        .await
29        .map_err(|e| format!("Failed to resolve tokens with RAG: {}", e))
30}
31
32/// Create an HTTP response with token resolution
33pub async fn create_token_resolved_response(
34    status: StatusCode,
35    body: Value,
36    use_rag: bool,
37    rag_config: Option<RagConfig>,
38) -> Response<Body> {
39    let resolved_body = if use_rag {
40        let config = rag_config.unwrap_or_default();
41        match resolve_response_tokens_with_rag(body.clone(), config).await {
42            Ok(resolved) => resolved,
43            Err(e) => {
44                error!(error = %e, "Failed to resolve tokens with RAG, using original body");
45                body
46            }
47        }
48    } else {
49        match resolve_response_tokens(body.clone()).await {
50            Ok(resolved) => resolved,
51            Err(e) => {
52                error!(error = %e, "Failed to resolve tokens, using original body");
53                body
54            }
55        }
56    };
57
58    let json_string = match serde_json::to_string_pretty(&resolved_body) {
59        Ok(s) => s,
60        Err(e) => {
61            error!(error = %e, "Failed to serialize response");
62            return (StatusCode::INTERNAL_SERVER_ERROR, "Failed to serialize response")
63                .into_response();
64        }
65    };
66
67    Response::builder()
68        .status(status)
69        .header("Content-Type", "application/json")
70        .body(Body::from(json_string))
71        .unwrap_or_else(|e| {
72            error!(error = %e, "Failed to build response");
73            (StatusCode::INTERNAL_SERVER_ERROR, "Failed to build response").into_response()
74        })
75}
76
77/// Token-resolved JSON response builder
78pub struct TokenResolvedResponse {
79    status: StatusCode,
80    body: Value,
81    use_rag: bool,
82    rag_config: Option<RagConfig>,
83}
84
85impl TokenResolvedResponse {
86    /// Create a new token-resolved response
87    pub fn new(status: StatusCode, body: Value) -> Self {
88        Self {
89            status,
90            body,
91            use_rag: false,
92            rag_config: None,
93        }
94    }
95
96    /// Enable RAG-based token resolution
97    pub fn with_rag(mut self, config: RagConfig) -> Self {
98        self.use_rag = true;
99        self.rag_config = Some(config);
100        self
101    }
102
103    /// Build the response
104    pub async fn build(self) -> Response<Body> {
105        create_token_resolved_response(self.status, self.body, self.use_rag, self.rag_config).await
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use mockforge_data::rag::LlmProvider;
113    use serde_json::json;
114
115    // ==================== Basic Token Resolution Tests ====================
116
117    #[tokio::test]
118    async fn test_resolve_response_tokens() {
119        let body = json!({
120            "id": "$random.uuid",
121            "name": "$faker.name",
122            "email": "$faker.email"
123        });
124
125        let result = resolve_response_tokens(body).await;
126        assert!(result.is_ok());
127
128        let resolved = result.unwrap();
129        assert!(resolved["id"].is_string());
130        assert!(resolved["name"].is_string());
131        assert!(resolved["email"].is_string());
132    }
133
134    #[tokio::test]
135    async fn test_resolve_nested_tokens() {
136        let body = json!({
137            "user": {
138                "id": "$random.uuid",
139                "profile": {
140                    "name": "$faker.name",
141                    "contact": {
142                        "email": "$faker.email",
143                        "phone": "$faker.phone"
144                    }
145                }
146            }
147        });
148
149        let result = resolve_response_tokens(body).await;
150        assert!(result.is_ok());
151
152        let resolved = result.unwrap();
153        assert!(resolved["user"]["id"].is_string());
154        assert!(resolved["user"]["profile"]["name"].is_string());
155        assert!(resolved["user"]["profile"]["contact"]["email"].is_string());
156    }
157
158    #[tokio::test]
159    async fn test_resolve_array_tokens() {
160        let body = json!({
161            "users": [
162                {"id": "$random.uuid", "name": "$faker.name"},
163                {"id": "$random.uuid", "name": "$faker.name"}
164            ]
165        });
166
167        let result = resolve_response_tokens(body).await;
168        assert!(result.is_ok());
169
170        let resolved = result.unwrap();
171        let users = resolved["users"].as_array().unwrap();
172        assert_eq!(users.len(), 2);
173        assert!(users[0]["id"].is_string());
174        assert!(users[0]["name"].is_string());
175    }
176
177    #[tokio::test]
178    async fn test_resolve_static_values() {
179        let body = json!({
180            "message": "Hello, World!",
181            "count": 42,
182            "active": true
183        });
184
185        let result = resolve_response_tokens(body.clone()).await;
186        assert!(result.is_ok());
187
188        let resolved = result.unwrap();
189        assert_eq!(resolved["message"], "Hello, World!");
190        assert_eq!(resolved["count"], 42);
191        assert_eq!(resolved["active"], true);
192    }
193
194    #[tokio::test]
195    async fn test_resolve_mixed_tokens_and_static() {
196        let body = json!({
197            "id": "$random.uuid",
198            "message": "Static message",
199            "count": 100
200        });
201
202        let result = resolve_response_tokens(body).await;
203        assert!(result.is_ok());
204
205        let resolved = result.unwrap();
206        assert!(resolved["id"].is_string());
207        assert_eq!(resolved["message"], "Static message");
208        assert_eq!(resolved["count"], 100);
209    }
210
211    #[tokio::test]
212    async fn test_resolve_empty_object() {
213        let body = json!({});
214        let result = resolve_response_tokens(body).await;
215        assert!(result.is_ok());
216        assert_eq!(result.unwrap(), json!({}));
217    }
218
219    #[tokio::test]
220    async fn test_resolve_null_value() {
221        let body = json!(null);
222        let result = resolve_response_tokens(body).await;
223        assert!(result.is_ok());
224    }
225
226    // ==================== TokenResolvedResponse Builder Tests ====================
227
228    #[tokio::test]
229    async fn test_token_resolved_response_builder() {
230        let body = json!({"message": "test"});
231        let response = TokenResolvedResponse::new(StatusCode::OK, body).build().await;
232
233        assert_eq!(response.status(), StatusCode::OK);
234    }
235
236    #[tokio::test]
237    async fn test_token_resolved_response_created() {
238        let body = json!({"id": "123", "created": true});
239        let response = TokenResolvedResponse::new(StatusCode::CREATED, body).build().await;
240
241        assert_eq!(response.status(), StatusCode::CREATED);
242    }
243
244    #[tokio::test]
245    async fn test_token_resolved_response_not_found() {
246        let body = json!({"error": "Resource not found"});
247        let response = TokenResolvedResponse::new(StatusCode::NOT_FOUND, body).build().await;
248
249        assert_eq!(response.status(), StatusCode::NOT_FOUND);
250    }
251
252    #[tokio::test]
253    async fn test_token_resolved_response_bad_request() {
254        let body = json!({"error": "Invalid input", "field": "email"});
255        let response = TokenResolvedResponse::new(StatusCode::BAD_REQUEST, body).build().await;
256
257        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
258    }
259
260    #[tokio::test]
261    async fn test_token_resolved_response_with_rag_config() {
262        let body = json!({"message": "test"});
263        let rag_config = RagConfig {
264            provider: LlmProvider::Ollama,
265            api_key: None,
266            model: "llama2".to_string(),
267            api_endpoint: "http://localhost:11434/api/generate".to_string(),
268            ..Default::default()
269        };
270
271        let builder = TokenResolvedResponse::new(StatusCode::OK, body).with_rag(rag_config);
272        assert!(builder.use_rag);
273        assert!(builder.rag_config.is_some());
274    }
275
276    #[test]
277    fn test_token_resolved_response_new_defaults() {
278        let body = json!({"test": "value"});
279        let response = TokenResolvedResponse::new(StatusCode::OK, body);
280
281        assert_eq!(response.status, StatusCode::OK);
282        assert!(!response.use_rag);
283        assert!(response.rag_config.is_none());
284    }
285
286    // ==================== create_token_resolved_response Tests ====================
287
288    #[tokio::test]
289    async fn test_create_response_ok() {
290        let body = json!({"status": "success"});
291        let response = create_token_resolved_response(StatusCode::OK, body, false, None).await;
292
293        assert_eq!(response.status(), StatusCode::OK);
294    }
295
296    #[tokio::test]
297    async fn test_create_response_with_tokens() {
298        let body = json!({
299            "id": "$random.uuid",
300            "timestamp": "$now"
301        });
302        let response = create_token_resolved_response(StatusCode::OK, body, false, None).await;
303
304        assert_eq!(response.status(), StatusCode::OK);
305    }
306
307    #[tokio::test]
308    async fn test_create_response_rag_disabled() {
309        let body = json!({"message": "test"});
310        let response = create_token_resolved_response(StatusCode::OK, body, false, None).await;
311
312        assert_eq!(response.status(), StatusCode::OK);
313    }
314
315    #[tokio::test]
316    async fn test_create_response_rag_enabled_no_config() {
317        // RAG enabled but no config provided - should use defaults
318        let body = json!({"message": "test"});
319        let response = create_token_resolved_response(StatusCode::OK, body, true, None).await;
320
321        // Should still succeed even if RAG fails (fallback to original body)
322        assert_eq!(response.status(), StatusCode::OK);
323    }
324
325    #[tokio::test]
326    async fn test_create_response_content_type_json() {
327        let body = json!({"data": "test"});
328        let response = create_token_resolved_response(StatusCode::OK, body, false, None).await;
329
330        let content_type = response.headers().get("Content-Type").and_then(|v| v.to_str().ok());
331        assert_eq!(content_type, Some("application/json"));
332    }
333
334    // ==================== RAG Config Tests ====================
335
336    #[test]
337    fn test_rag_config_default() {
338        let config = RagConfig::default();
339        assert!(config.temperature >= 0.0);
340        assert!(config.max_tokens > 0);
341    }
342
343    #[test]
344    fn test_rag_config_with_provider() {
345        let config = RagConfig {
346            provider: LlmProvider::OpenAI,
347            api_key: Some("test-key".to_string()),
348            model: "gpt-4".to_string(),
349            api_endpoint: "https://api.openai.com/v1/chat/completions".to_string(),
350            ..Default::default()
351        };
352
353        assert!(matches!(config.provider, LlmProvider::OpenAI));
354        assert_eq!(config.api_key, Some("test-key".to_string()));
355    }
356
357    // ==================== Edge Cases ====================
358
359    #[tokio::test]
360    async fn test_resolve_deeply_nested_array() {
361        let body = json!({
362            "data": {
363                "items": [
364                    {"ids": ["$random.uuid", "$random.uuid"]},
365                    {"ids": ["$random.uuid"]}
366                ]
367            }
368        });
369
370        let result = resolve_response_tokens(body).await;
371        assert!(result.is_ok());
372    }
373
374    #[tokio::test]
375    async fn test_resolve_complex_structure() {
376        let body = json!({
377            "meta": {
378                "total": 100,
379                "page": 1
380            },
381            "data": [
382                {
383                    "id": "$random.uuid",
384                    "attributes": {
385                        "name": "$faker.name",
386                        "created_at": "$now"
387                    },
388                    "relationships": {
389                        "author": {
390                            "id": "$random.uuid"
391                        }
392                    }
393                }
394            ]
395        });
396
397        let result = resolve_response_tokens(body).await;
398        assert!(result.is_ok());
399    }
400}
mockforge_http/token_response.rs

mockforge_http/
token_response.rs
