Skip to main content

kftray_http_logs/
message.rs

1use bytes::Bytes;
2use chrono::{
3    DateTime,
4    Utc,
5};
6
7use crate::formatter::MessageFormatter;
8
9#[derive(Debug, Clone)]
10pub enum HttpMessage {
11    Request {
12        trace_id: String,
13        timestamp: DateTime<Utc>,
14        buffer: Bytes,
15    },
16    Response {
17        trace_id: String,
18        timestamp: DateTime<Utc>,
19        took_ms: i64,
20        buffer: Bytes,
21    },
22}
23
24#[derive(Debug, Clone)]
25pub enum LogMessage {
26    Request(String),
27    Response(String),
28    PreformattedResponse(String),
29    TriggerFlush,
30}
31
32impl LogMessage {
33    pub fn as_bytes(&self) -> &[u8] {
34        match self {
35            LogMessage::Request(log) => log.as_bytes(),
36            LogMessage::Response(log) => log.as_bytes(),
37            LogMessage::PreformattedResponse(log) => log.as_bytes(),
38            LogMessage::TriggerFlush => &[],
39        }
40    }
41
42    pub fn size(&self) -> usize {
43        match self {
44            LogMessage::TriggerFlush => 0,
45            _ => self.as_bytes().len(),
46        }
47    }
48
49    pub fn message_type(&self) -> &'static str {
50        match self {
51            LogMessage::Request(_) => "Request",
52            LogMessage::Response(_) => "Response",
53            LogMessage::PreformattedResponse(_) => "PreformattedResponse",
54            LogMessage::TriggerFlush => "TriggerFlush",
55        }
56    }
57
58    pub fn is_response(&self) -> bool {
59        matches!(
60            self,
61            LogMessage::Response(_) | LogMessage::PreformattedResponse(_)
62        )
63    }
64
65    pub fn is_flush_trigger(&self) -> bool {
66        matches!(self, LogMessage::TriggerFlush)
67    }
68
69    pub fn new_preformatted_response(
70        trace_id: String, timestamp: DateTime<Utc>, took_ms: i64, buffer: Bytes,
71    ) -> Self {
72        let formatted =
73            MessageFormatter::format_preformatted_response(&trace_id, timestamp, took_ms, &buffer);
74
75        LogMessage::PreformattedResponse(formatted)
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use bytes::Bytes;
82    use chrono::Utc;
83
84    use super::*;
85
86    #[test]
87    fn test_http_message_request() {
88        let trace_id = "test-trace-123".to_string();
89        let timestamp = Utc::now();
90        let buffer = Bytes::from("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n");
91
92        let msg = HttpMessage::Request {
93            trace_id: trace_id.clone(),
94            timestamp,
95            buffer: buffer.clone(),
96        };
97
98        if let HttpMessage::Request {
99            trace_id: tid,
100            timestamp: ts,
101            buffer: buf,
102        } = msg
103        {
104            assert_eq!(tid, trace_id);
105            assert_eq!(ts, timestamp);
106            assert_eq!(buf, buffer);
107        } else {
108            panic!("Expected HttpMessage::Request");
109        }
110    }
111
112    #[test]
113    fn test_http_message_response() {
114        let trace_id = "test-trace-456".to_string();
115        let timestamp = Utc::now();
116        let took_ms = 123;
117        let buffer = Bytes::from("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nHello");
118
119        let msg = HttpMessage::Response {
120            trace_id: trace_id.clone(),
121            timestamp,
122            took_ms,
123            buffer: buffer.clone(),
124        };
125
126        if let HttpMessage::Response {
127            trace_id: tid,
128            timestamp: ts,
129            took_ms: tm,
130            buffer: buf,
131        } = msg
132        {
133            assert_eq!(tid, trace_id);
134            assert_eq!(ts, timestamp);
135            assert_eq!(tm, took_ms);
136            assert_eq!(buf, buffer);
137        } else {
138            panic!("Expected HttpMessage::Response");
139        }
140    }
141
142    #[test]
143    fn test_log_message_as_bytes() {
144        let req_content = "Request content".to_string();
145        let resp_content = "Response content".to_string();
146        let preformatted = "Preformatted content".to_string();
147
148        let req_msg = LogMessage::Request(req_content.clone());
149        let resp_msg = LogMessage::Response(resp_content.clone());
150        let preformatted_msg = LogMessage::PreformattedResponse(preformatted.clone());
151        let flush_msg = LogMessage::TriggerFlush;
152
153        assert_eq!(req_msg.as_bytes(), req_content.as_bytes());
154        assert_eq!(resp_msg.as_bytes(), resp_content.as_bytes());
155        assert_eq!(preformatted_msg.as_bytes(), preformatted.as_bytes());
156
157        // Use Vec type explicitly to fix the error
158        let empty_vec: Vec<u8> = Vec::new();
159        assert_eq!(flush_msg.as_bytes(), empty_vec.as_slice());
160    }
161
162    #[test]
163    fn test_log_message_size() {
164        let req_content = "Request content".to_string();
165        let resp_content = "Response content".to_string();
166        let preformatted = "Preformatted content".to_string();
167
168        let req_msg = LogMessage::Request(req_content.clone());
169        let resp_msg = LogMessage::Response(resp_content.clone());
170        let preformatted_msg = LogMessage::PreformattedResponse(preformatted.clone());
171        let flush_msg = LogMessage::TriggerFlush;
172
173        assert_eq!(req_msg.size(), req_content.len());
174        assert_eq!(resp_msg.size(), resp_content.len());
175        assert_eq!(preformatted_msg.size(), preformatted.len());
176        assert_eq!(flush_msg.size(), 0);
177    }
178
179    #[test]
180    fn test_log_message_type() {
181        let req_msg = LogMessage::Request("req".to_string());
182        let resp_msg = LogMessage::Response("resp".to_string());
183        let preformatted_msg = LogMessage::PreformattedResponse("pre".to_string());
184        let flush_msg = LogMessage::TriggerFlush;
185
186        assert_eq!(req_msg.message_type(), "Request");
187        assert_eq!(resp_msg.message_type(), "Response");
188        assert_eq!(preformatted_msg.message_type(), "PreformattedResponse");
189        assert_eq!(flush_msg.message_type(), "TriggerFlush");
190    }
191
192    #[test]
193    fn test_is_response() {
194        let req_msg = LogMessage::Request("req".to_string());
195        let resp_msg = LogMessage::Response("resp".to_string());
196        let preformatted_msg = LogMessage::PreformattedResponse("pre".to_string());
197        let flush_msg = LogMessage::TriggerFlush;
198
199        assert!(!req_msg.is_response());
200        assert!(resp_msg.is_response());
201        assert!(preformatted_msg.is_response());
202        assert!(!flush_msg.is_response());
203    }
204
205    #[test]
206    fn test_is_flush_trigger() {
207        let req_msg = LogMessage::Request("req".to_string());
208        let resp_msg = LogMessage::Response("resp".to_string());
209        let preformatted_msg = LogMessage::PreformattedResponse("pre".to_string());
210        let flush_msg = LogMessage::TriggerFlush;
211
212        assert!(!req_msg.is_flush_trigger());
213        assert!(!resp_msg.is_flush_trigger());
214        assert!(!preformatted_msg.is_flush_trigger());
215        assert!(flush_msg.is_flush_trigger());
216    }
217
218    #[test]
219    fn test_new_preformatted_response() {
220        let formatted = "Test formatted response".to_string();
221        let msg = LogMessage::PreformattedResponse(formatted);
222
223        if let LogMessage::PreformattedResponse(content) = msg {
224            assert_eq!(content, "Test formatted response");
225        } else {
226            panic!("Expected LogMessage::PreformattedResponse");
227        }
228    }
229}