union_square 0.2.0

A proxy/wire-tap service for making LLM calls and recording everything that happens in a session for later analysis and test-case extraction
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259

//! Test helpers specific to middleware testing
//!
//! This module provides utilities for testing individual middleware
//! components in isolation.

#[cfg(test)]
pub mod helpers {
    use crate::proxy::{AuthConfig, ProxyMiddlewareStack};
    use axum::{
        body::Body,
        http::{Request, Response, StatusCode},
        Router,
    };
    use tower::ServiceExt;

    /// Test harness for middleware testing
    pub struct MiddlewareTestHarness {
        router: Router,
    }

    impl Default for MiddlewareTestHarness {
        fn default() -> Self {
            Self::new()
        }
    }

    impl MiddlewareTestHarness {
        /// Create a new test harness with a simple handler
        pub fn new() -> Self {
            let router = Router::new()
                .route("/test", axum::routing::any(test_handler))
                .route("/echo", axum::routing::any(echo_handler))
                .route("/error", axum::routing::any(error_handler));

            Self { router }
        }

        /// Apply a single middleware for testing
        /// Note: This is a placeholder - actual middleware application
        /// should use the ProxyMiddlewareStack for proper composition
        pub fn with_middleware(self) -> Self {
            // Middleware is applied via with_full_stack method
            self
        }

        /// Apply the full middleware stack
        pub fn with_full_stack(self, auth_config: AuthConfig) -> Self {
            let stack = ProxyMiddlewareStack::new(auth_config);
            Self {
                router: stack.apply_to_router(self.router),
            }
        }

        /// Send a request through the middleware stack
        pub async fn send_request(&mut self, request: Request<Body>) -> Response<Body> {
            self.router
                .clone()
                .oneshot(request)
                .await
                .expect("middleware should not fail")
        }

        /// Helper to create and send a GET request
        pub async fn get(&mut self, path: &str) -> Response<Body> {
            let request = Request::builder()
                .method("GET")
                .uri(path)
                .body(Body::empty())
                .expect("test request should be valid");
            self.send_request(request).await
        }

        /// Helper to create and send an authenticated request
        pub async fn get_with_auth(&mut self, path: &str, api_key: &str) -> Response<Body> {
            let request = Request::builder()
                .method("GET")
                .uri(path)
                .header("Authorization", format!("Bearer {api_key}"))
                .body(Body::empty())
                .expect("test request should be valid");
            self.send_request(request).await
        }
    }

    /// Simple test handler that returns OK
    async fn test_handler() -> StatusCode {
        StatusCode::OK
    }

    /// Echo handler that returns request details
    async fn echo_handler(req: Request<Body>) -> Response<Body> {
        let headers = req
            .headers()
            .iter()
            .map(|(k, v)| format!("{}: {}", k, v.to_str().unwrap_or("<binary>")))
            .collect::<Vec<_>>()
            .join("\n");

        Response::builder()
            .status(StatusCode::OK)
            .body(Body::from(headers))
            .expect("response should be valid")
    }

    /// Error handler that always returns an error
    async fn error_handler() -> Result<StatusCode, TestError> {
        Err(TestError::Simulated)
    }

    #[derive(Debug)]
    enum TestError {
        Simulated,
    }

    impl axum::response::IntoResponse for TestError {
        fn into_response(self) -> axum::response::Response {
            (StatusCode::INTERNAL_SERVER_ERROR, "Test error").into_response()
        }
    }

    /// Assertions for middleware testing
    pub mod assertions {
        use super::*;

        /// Assert that a response has the expected status code
        pub fn assert_status(response: &Response<Body>, expected: StatusCode) {
            assert_eq!(
                response.status(),
                expected,
                "Expected status {} but got {}",
                expected,
                response.status()
            );
        }

        /// Assert that a response contains a specific header
        pub fn assert_header_exists(response: &Response<Body>, header_name: &str) {
            assert!(
                response.headers().contains_key(header_name),
                "Expected header '{header_name}' not found"
            );
        }

        /// Assert that a response header has a specific value
        pub fn assert_header_value(
            response: &Response<Body>,
            header_name: &str,
            expected_value: &str,
        ) {
            let actual = response
                .headers()
                .get(header_name)
                .and_then(|v| v.to_str().ok())
                .unwrap_or("<missing>");

            assert_eq!(
                actual, expected_value,
                "Header '{header_name}' expected '{expected_value}' but got '{actual}'"
            );
        }

        /// Assert that response body contains expected text
        pub async fn assert_body_contains(response: Response<Body>, expected: &str) {
            let body_bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
                .await
                .expect("body should be readable");
            let body_str = String::from_utf8_lossy(&body_bytes);

            assert!(
                body_str.contains(expected),
                "Body does not contain '{expected}'. Actual: '{body_str}'"
            );
        }
    }

    /// Middleware test scenarios
    pub mod scenarios {
        use super::*;
        use crate::proxy::headers::X_REQUEST_ID;

        /// Test that request ID middleware adds IDs to requests
        pub async fn test_request_id_generation(harness: &mut MiddlewareTestHarness) {
            let response = harness.get("/test").await;
            assertions::assert_header_exists(&response, X_REQUEST_ID);

            // Verify UUID format
            let request_id = response
                .headers()
                .get(X_REQUEST_ID)
                .and_then(|v| v.to_str().ok())
                .expect("request ID should exist");

            assert!(
                uuid::Uuid::parse_str(request_id).is_ok(),
                "Request ID should be a valid UUID"
            );
        }

        /// Test that existing request IDs are preserved
        pub async fn test_request_id_preservation(harness: &mut MiddlewareTestHarness) {
            let existing_id = uuid::Uuid::new_v4().to_string();
            let request = Request::builder()
                .uri("/test")
                .header(X_REQUEST_ID, &existing_id)
                .body(Body::empty())
                .expect("request should be valid");

            let response = harness.send_request(request).await;
            assertions::assert_header_value(&response, X_REQUEST_ID, &existing_id);
        }

        /// Test authentication with valid API key
        pub async fn test_valid_authentication(
            harness: &mut MiddlewareTestHarness,
            valid_key: &str,
        ) {
            let response = harness.get_with_auth("/test", valid_key).await;
            assertions::assert_status(&response, StatusCode::OK);
        }

        /// Test authentication with invalid API key
        pub async fn test_invalid_authentication(harness: &mut MiddlewareTestHarness) {
            let response = harness.get_with_auth("/test", "invalid-key").await;
            assertions::assert_status(&response, StatusCode::UNAUTHORIZED);
        }

        /// Test health check bypass
        pub async fn test_health_check_bypass(harness: &mut MiddlewareTestHarness) {
            let response = harness.get("/health").await;
            assertions::assert_status(&response, StatusCode::OK);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::helpers::*;
    use axum::http::StatusCode;

    #[tokio::test]
    async fn test_middleware_harness() {
        let mut harness = MiddlewareTestHarness::new();
        let response = harness.get("/test").await;
        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_echo_handler() {
        let mut harness = MiddlewareTestHarness::new();
        let response = harness.get_with_auth("/echo", "test-key").await;

        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_str = String::from_utf8_lossy(&body);

        assert!(body_str.contains("authorization: Bearer test-key"));
    }
}