Add response status to AWS and remote namespaces

README.md

@@ -8,7 +8,7 @@ Add the following to your `Cargo.toml` file:
 
 ```toml
 [dependencies]
-xray-lite = { git = "https://github.com/codemonger-io/xray-lite.git", tag = "v0.0.5" }
+xray-lite = { git = "https://github.com/codemonger-io/xray-lite.git", tag = "v0.0.6" }
 ```
 
 ## Usage

xray-lite/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "xray-lite"
-version = "0.0.5"
+version = "0.0.6"
 authors = ["softprops <[email protected]>", "Kikuo Emoto <[email protected]>"]
 edition = "2021"
 description = "AWS X-Ray daemon client for Rust"

xray-lite/src/namespace.rs

@@ -1,6 +1,6 @@
 //! Namespace encapsulation for subsegments.
 
-use crate::segment::{AwsOperation, Http, Request, Subsegment};
+use crate::segment::{AwsOperation, Http, Request, Response, Subsegment};
 
 /// Namespace.
 pub trait Namespace {
@@ -19,6 +19,7 @@ pub struct AwsNamespace {
     service: String,
     operation: String,
     request_id: Option<String>,
+    response_status: Option<u16>,
 }
 
 impl AwsNamespace {
@@ -28,6 +29,7 @@ impl AwsNamespace {
             service: service.into(),
             operation: operation.into(),
             request_id: None,
+            response_status: None,
         }
     }
 
@@ -36,6 +38,12 @@ impl AwsNamespace {
         self.request_id = Some(request_id.into());
         self
     }
+
+    /// Sets the response status.
+    pub fn response_status(&mut self, status: u16) -> &mut Self {
+        self.response_status = Some(status);
+        self
+    }
 }
 
 impl Namespace for AwsNamespace {
@@ -61,6 +69,28 @@ impl Namespace for AwsNamespace {
                 ..AwsOperation::default()
             });
         }
+        if let Some(response_status) = self.response_status {
+            if let Some(http) = subsegment.http.as_mut() {
+                if let Some(response) = http.response.as_mut() {
+                    if response.status.is_none() {
+                        response.status = Some(response_status);
+                    }
+                } else {
+                    http.response = Some(Response {
+                        status: Some(response_status),
+                        ..Response::default()
+                    });
+                }
+            } else {
+                subsegment.http = Some(Http {
+                    response: Some(Response {
+                        status: Some(response_status),
+                        ..Response::default()
+                    }),
+                    ..Http::default()
+                });
+            }
+        }
     }
 }
 
@@ -70,6 +100,7 @@ pub struct RemoteNamespace {
     name: String,
     method: String,
     url: String,
+    response_status: Option<u16>,
 }
 
 impl RemoteNamespace {
@@ -79,8 +110,15 @@ impl RemoteNamespace {
             name: name.into(),
             method: method.into(),
             url: url.into(),
+            response_status: None,
         }
     }
+
+    /// Sets the response status.
+    pub fn response_status(&mut self, status: u16) -> &mut Self {
+        self.response_status = Some(status);
+        self
+    }
 }
 
 impl Namespace for RemoteNamespace {
@@ -117,6 +155,19 @@ impl Namespace for RemoteNamespace {
                 ..Http::default()
             });
         }
+        if let Some(response_status) = self.response_status {
+            let http = subsegment.http.as_mut().expect("http must have been set");
+            if let Some(response) = http.response.as_mut() {
+                if response.status.is_none() {
+                    response.status = Some(response_status);
+                }
+            } else {
+                http.response = Some(Response {
+                    status: Some(response_status),
+                    ..Response::default()
+                });
+            }
+        }
     }
 }
 
@@ -141,3 +192,94 @@ impl Namespace for CustomNamespace {
     // does nothing
     fn update_subsegment(&self, _subsegment: &mut Subsegment) {}
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn aws_namespace_should_have_service_name_as_name() {
+        let namespace = AwsNamespace::new("S3", "GetObject");
+        assert_eq!(namespace.name(""), "S3");
+        assert_eq!(namespace.name("prefix"), "S3");
+    }
+
+    #[test]
+    fn aws_namespace_should_update_subsegment_with_aws_operation() {
+        let namespace = AwsNamespace::new("S3", "GetObject");
+        let mut subsegment = Subsegment::default();
+        namespace.update_subsegment(&mut subsegment);
+        assert_eq!(subsegment.namespace.unwrap(), "aws");
+        assert_eq!(subsegment.aws.unwrap().operation.unwrap(), "GetObject");
+    }
+
+    #[test]
+    fn aws_namespace_should_update_subsegment_with_request_id() {
+        let mut namespace = AwsNamespace::new("S3", "GetObject");
+        namespace.request_id("12345");
+        let mut subsegment = Subsegment::default();
+        namespace.update_subsegment(&mut subsegment);
+        assert_eq!(subsegment.aws.unwrap().request_id.unwrap(), "12345");
+    }
+
+    #[test]
+    fn aws_namespace_should_update_subsegment_with_response_status() {
+        let mut namespace = AwsNamespace::new("S3", "GetObject");
+        namespace.response_status(200);
+        let mut subsegment = Subsegment::default();
+        namespace.update_subsegment(&mut subsegment);
+        assert_eq!(
+            subsegment
+                .http
+                .expect("http")
+                .response
+                .expect("response")
+                .status
+                .expect("status"),
+            200,
+        );
+    }
+
+    #[test]
+    fn remote_namespace_should_have_name_as_name() {
+        let namespace = RemoteNamespace::new("codemonger.io", "GET", "https://codemonger.io/");
+        assert_eq!(namespace.name(""), "codemonger.io");
+        assert_eq!(namespace.name("prefix"), "codemonger.io");
+    }
+
+    #[test]
+    fn remote_namespace_should_update_subsegment_with_remote_service() {
+        let namespace = RemoteNamespace::new("codemonger.io", "GET", "https://codemonger.io/");
+        let mut subsegment = Subsegment::default();
+        namespace.update_subsegment(&mut subsegment);
+        assert_eq!(subsegment.namespace.unwrap(), "remote");
+        let request = subsegment.http.expect("http").request.expect("request");
+        assert_eq!(request.method.unwrap(), "GET");
+        assert_eq!(request.url.unwrap(), "https://codemonger.io/");
+    }
+
+    #[test]
+    fn remote_namespace_should_update_subsegment_with_response_status() {
+        let mut namespace = RemoteNamespace::new("codemonger.io", "GET", "https://codemonger.io/");
+        namespace.response_status(200);
+        let mut subsegment = Subsegment::default();
+        namespace.update_subsegment(&mut subsegment);
+        assert_eq!(
+            subsegment
+                .http
+                .expect("http")
+                .response
+                .expect("response")
+                .status
+                .expect("status"),
+            200,
+        );
+    }
+
+    #[test]
+    fn custom_namespace_should_have_prefixed_name() {
+        let namespace = CustomNamespace::new("TestSubsegment");
+        assert_eq!(namespace.name(""), "TestSubsegment");
+        assert_eq!(namespace.name("prefix"), "prefixTestSubsegment");
+    }
+}