add: new tests for to_bytes and from_bytes and changed name

niclabs · Jan 6, 2025 · 44aa939 · 44aa939
1 parent bcfb9fc
commit 44aa939
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 3 deletions.
diff --git a/src/message/rdata/opt_rdata.rs b/src/message/rdata/opt_rdata.rs
@@ -75,7 +75,7 @@ impl FromBytes<Result<Self, &'static str>> for OptRdata {
 
             i += option_length as usize;
 
-            let option_data = OptionData::from_with_opt_type(option_data, option_code);
+            let option_data = OptionData::from_bytes_with_opt_type(option_data, option_code);
 
             if let Err(_) = option_data {
                 return Err("Format Error");

diff --git a/src/message/rdata/opt_rdata/option_data.rs b/src/message/rdata/opt_rdata/option_data.rs
@@ -45,7 +45,7 @@ impl ToBytes for OptionData {
 }
 
 impl OptionData {
-    pub fn from_with_opt_type(bytes: Vec<u8>, opt_t: OptionCode) -> Result<OptionData, &'static str> {
+    pub fn from_bytes_with_opt_type(bytes: Vec<u8>, opt_t: OptionCode) -> Result<OptionData, &'static str> {
         match opt_t {
             OptionCode::NSID => {
                 let nsid = String::from_utf8(bytes).map_err(|_| "Error parsing NSID")?;
@@ -61,4 +61,63 @@ impl OptionData {
             _ => Ok(OptionData::Unknown(bytes))
         }
     }
-}
+}
+
+#[cfg(test)]
+mod option_data_tests {
+    use super::*;
+    use crate::message::rdata::opt_rdata::ede_optdata::EdeOptData;
+    use crate::message::rdata::opt_rdata::ede_code::EdeCode;
+    use crate::message::rdata::opt_rdata::option_code::OptionCode;
+
+    #[test]
+    fn test_option_data_nsid() {
+        let nsid_string = "testNSID".to_string();
+        let option_data = OptionData::NSID(nsid_string.clone());
+        let serialized = option_data.to_bytes();
+        let rebuilt = OptionData::from_bytes_with_opt_type(serialized.clone(), OptionCode::NSID)
+            .expect("NSID reconstruction failed");
+        assert_eq!(option_data, rebuilt);
+        assert_eq!(serialized.len(), nsid_string.len());
+    }
+
+    #[test]
+    fn test_option_data_ede() {
+        let code = EdeCode::StaleAns;
+        let msg = "Stale Answer".to_string();
+        let ede_data = EdeOptData::new(code, msg.clone());
+        let option_data = OptionData::EDE(ede_data.clone());
+        let serialized = option_data.to_bytes();
+        let rebuilt = OptionData::from_bytes_with_opt_type(serialized.clone(), OptionCode::EDE)
+            .expect("EDE reconstruction failed");
+        assert_eq!(option_data, rebuilt);
+        if let OptionData::EDE(ede_rebuilt) = rebuilt {
+            assert_eq!(ede_rebuilt.get_info_code(), code);
+            assert_eq!(ede_rebuilt.get_extra_text(), msg);
+        } else {
+            panic!("Expected OptionData::EDE variant");
+        }
+    }
+
+    #[test]
+    fn test_option_data_padding() {
+        let padding_bytes = vec![0x00, 0x00, 0x00, 0x00];
+        let option_data = OptionData::Padding(padding_bytes.clone());
+        let serialized = option_data.to_bytes();
+        let rebuilt = OptionData::from_bytes_with_opt_type(serialized.clone(), OptionCode::PADDING)
+            .expect("Padding reconstruction failed");
+        assert_eq!(option_data, rebuilt);
+        assert_eq!(serialized.len(), 4);
+    }
+
+    #[test]
+    fn test_option_data_unknown() {
+        let unknown_bytes = vec![0xde, 0xad, 0xbe, 0xef];
+        let option_data = OptionData::Unknown(unknown_bytes.clone());
+        let serialized = option_data.to_bytes();
+        let rebuilt = OptionData::from_bytes_with_opt_type(serialized.clone(), OptionCode::UNKNOWN(999))
+            .expect("Unknown reconstruction failed");
+        assert_eq!(option_data, rebuilt);
+        assert_eq!(serialized.len(), 4);
+    }
+}