From e673db3de2d7fd8ae87edaf1da8fe21c579643a2 Mon Sep 17 00:00:00 2001
From: Abdulla Abdurakhmanov <me@abdolence.dev>
Date: Fri, 14 Jul 2023 12:17:58 +0200
Subject: [PATCH] Fixed missing content type and added headers() function to
 override from apps

---
 examples/json-example.rs |  2 +-
 src/csv_format.rs        |  6 ++++++
 src/json_formats.rs      | 14 ++++++++++++++
 src/protobuf_format.rs   |  6 ++++++
 src/stream_body_as.rs    | 17 +++++++++++++++--
 5 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/examples/json-example.rs b/examples/json-example.rs
index ae3de78..6f0d398 100644
--- a/examples/json-example.rs
+++ b/examples/json-example.rs
@@ -20,7 +20,7 @@ fn source_test_stream() -> impl Stream<Item = MyTestStructure> {
         MyTestStructure {
             some_test_field: "test1".to_string()
         };
-        1000
+        100
     ])
     .throttle(std::time::Duration::from_millis(50))
 }
diff --git a/src/csv_format.rs b/src/csv_format.rs
index 16d01c4..9ff4810 100644
--- a/src/csv_format.rs
+++ b/src/csv_format.rs
@@ -116,6 +116,12 @@ mod tests {
             .add("\n");
 
         let res = client.get("/").send().await.unwrap();
+        assert_eq!(
+            res.headers()
+                .get("content-type")
+                .and_then(|h| h.to_str().ok()),
+            Some("text/csv")
+        );
         let body = res.text().await.unwrap();
 
         assert_eq!(body, expected_csv);
diff --git a/src/json_formats.rs b/src/json_formats.rs
index 548e1ff..3193e8b 100644
--- a/src/json_formats.rs
+++ b/src/json_formats.rs
@@ -158,6 +158,13 @@ mod tests {
 
         let expected_json = serde_json::to_string(&test_stream_vec).unwrap();
         let res = client.get("/").send().await.unwrap();
+        assert_eq!(
+            res.headers()
+                .get("content-type")
+                .and_then(|h| h.to_str().ok()),
+            Some("application/json")
+        );
+
         let body = res.text().await.unwrap();
 
         assert_eq!(body, expected_json);
@@ -193,6 +200,13 @@ mod tests {
             .join("\n");
 
         let res = client.get("/").send().await.unwrap();
+        assert_eq!(
+            res.headers()
+                .get("content-type")
+                .and_then(|h| h.to_str().ok()),
+            Some("application/jsonstream")
+        );
+
         let body = res.text().await.unwrap();
 
         assert_eq!(body, expected_json);
diff --git a/src/protobuf_format.rs b/src/protobuf_format.rs
index c135335..e323033 100644
--- a/src/protobuf_format.rs
+++ b/src/protobuf_format.rs
@@ -112,6 +112,12 @@ mod tests {
             .collect();
 
         let res = client.get("/").send().await.unwrap();
+        assert_eq!(
+            res.headers()
+                .get("content-type")
+                .and_then(|h| h.to_str().ok()),
+            Some("application/x-protobuf-stream")
+        );
         let body = res.bytes().await.unwrap().to_vec();
 
         assert_eq!(body, expected_proto_buf);
diff --git a/src/stream_body_as.rs b/src/stream_body_as.rs
index aacd3df..a8632e7 100644
--- a/src/stream_body_as.rs
+++ b/src/stream_body_as.rs
@@ -31,11 +31,24 @@ impl<'a> StreamBodyAs<'a> {
             trailers: stream_format.http_response_trailers(),
         }
     }
+
+    pub fn headers(mut self, headers: HeaderMap) -> Self {
+        self.trailers = Some(headers);
+        self
+    }
 }
 
 impl IntoResponse for StreamBodyAs<'static> {
-    fn into_response(self) -> Response {
-        Response::new(axum::body::boxed(self))
+    fn into_response(mut self) -> Response {
+        let headers = if let Some(trailers) = self.trailers.take() {
+            trailers
+        } else {
+            HeaderMap::new()
+        };
+
+        let mut response = Response::new(axum::body::boxed(self));
+        *response.headers_mut() = headers;
+        response
     }
 }