Reduce test verbosity in arp

In arp, there was originally lots of repetitive code. This was removed. Signed-off-by: Andrew Yao <[email protected]>
firecracker-microvm · May 7, 2024 · dfed9e1 · dfed9e1
1 parent 81e8b68
commit dfed9e1
Showing 1 changed file with 63 additions and 96 deletions.
diff --git a/src/vmm/src/dumbo/pdu/arp.rs b/src/vmm/src/dumbo/pdu/arp.rs
@@ -403,102 +403,69 @@ mod tests {
             ArpError::Operation
         );
 
-        // TODO: The following test code is way more verbose than it should've been. Make it
-        // prettier at some point.
-
-        // Let's write a valid request.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap();
-
-        // Now we start writing invalid requests. We've already tried with an invalid operation.
-
-        // Invalid htype.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET + 1,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::HType
-        );
-
-        // Invalid ptype.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4 + 1,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::PType
-        );
-
-        // Invalid hlen.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN + 1,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::HLen
-        );
-
-        // Invalid plen.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN + 1,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::PLen
-        );
+        // Various requests
+        let requests = [
+            (
+                HTYPE_ETHERNET,
+                ETHERTYPE_IPV4,
+                MAC_ADDR_LEN,
+                IPV4_ADDR_LEN,
+                None,
+            ), // Valid request
+            (
+                HTYPE_ETHERNET + 1,
+                ETHERTYPE_IPV4,
+                MAC_ADDR_LEN,
+                IPV4_ADDR_LEN,
+                Some(ArpError::HType),
+            ), // Invalid htype
+            (
+                HTYPE_ETHERNET,
+                ETHERTYPE_IPV4 + 1,
+                MAC_ADDR_LEN,
+                IPV4_ADDR_LEN,
+                Some(ArpError::PType),
+            ), // Invalid ptype
+            (
+                HTYPE_ETHERNET,
+                ETHERTYPE_IPV4,
+                MAC_ADDR_LEN + 1,
+                IPV4_ADDR_LEN,
+                Some(ArpError::HLen),
+            ), // Invalid hlen
+            (
+                HTYPE_ETHERNET,
+                ETHERTYPE_IPV4,
+                MAC_ADDR_LEN,
+                IPV4_ADDR_LEN + 1,
+                Some(ArpError::PLen),
+            ), // Invalid plen
+        ];
+
+        for (htype, ptype, hlen, plen, err) in requests.iter() {
+            EthIPv4ArpFrame::write_raw(
+                &mut a[..ETH_IPV4_FRAME_LEN],
+                *htype,
+                *ptype,
+                *hlen,
+                *plen,
+                OPER_REQUEST,
+                sha,
+                spa,
+                tha,
+                tpa,
+            )
+            .unwrap();
+            match err {
+                None => {
+                    EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap();
+                }
+                Some(arp_error) => assert_eq!(
+                    EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
+                    *arp_error
+                ),
+            }
+        }
     }
 
     #[test]