Added headers to source message

Signed-off-by: shubham <[email protected]>
numaproj · May 28, 2024 · 085d3a6 · 085d3a6
1 parent 2483003
commit 085d3a6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/examples/simple-source/src/main.rs b/examples/simple-source/src/main.rs
@@ -53,7 +53,8 @@ pub(crate) mod simple_source {
                 self.read_idx
                     .store(self.read_idx.load(Ordering::Relaxed) + 1, Ordering::Relaxed);
                 let offset = self.read_idx.load(Ordering::Relaxed);
-
+                let mut headers=HashMap::new();
+                headers.insert(String::from("key"),String::from("key"));
                 // send the message to the transmitter
                 transmitter
                     .send(Message {
@@ -64,7 +65,7 @@ pub(crate) mod simple_source {
                         },
                         event_time: chrono::offset::Utc::now(),
                         keys: vec![],
-                        headers:HashMap::new().insert("key","key")
+                        headers:headers.clone()
                     })
                     .await
                     .unwrap();

diff --git a/src/source.rs b/src/source.rs
@@ -291,7 +291,7 @@ impl<T> Server<T> {
 mod tests {
     use super::proto;
     use chrono::Utc;
-    use std::collections::HashSet;
+    use std::collections::{HashMap, HashSet};
     use std::vec;
     use std::{error::Error, time::Duration};
     use tokio_stream::StreamExt;
@@ -323,6 +323,8 @@ mod tests {
         async fn read(&self, request: SourceReadRequest, transmitter: Sender<Message>) {
             let event_time = Utc::now();
             let mut message_offsets = Vec::with_capacity(request.count);
+            let mut headers=HashMap::new();
+            headers.insert(String::from("key"),String::from("key"));
             for i in 0..request.count {
                 // we assume timestamp in nanoseconds would be unique on each read operation from our source
                 let offset = format!("{}-{}", event_time.timestamp_nanos_opt().unwrap(), i);
@@ -335,6 +337,7 @@ mod tests {
                             partition_id: 0,
                         },
                         keys: vec![],
+                        headers:headers.clone(),
                     })
                     .await
                     .unwrap();