kftray_http_logs/
message.rs1use 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 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}