test_conn.py: update unit tests

stratosphereips · Jan 15, 2025 · c3424d4 · c3424d4
1 parent aa7500b
commit c3424d4
Showing 1 changed file with 36 additions and 12 deletions.
diff --git a/tests/test_conn.py b/tests/test_conn.py
@@ -451,21 +451,43 @@ def test_is_interface_timeout_reached(mock_time_diff, expected_result):
 
 @pytest.mark.parametrize(
     "flow_type, appproto, daddr, input_type, "
-    "is_doh_server, is_dns_server, "
+    "is_doh_server, is_dns_server, is_dhcp_server,"
     "client_ips, expected_result",
     [
-        # # Testcase 1: Not a 'conn' flow type
-        # ("dns", "dns", "8.8.8.8", "pcap", False, False, [], True),
-        # # Testcase 2: DNS application protocol
-        # ("conn", "dns", "8.8.8.8", "pcap", False, False, [], True),
-        # # Testcase 3: Ignored IP
-        # ("conn", "http", "192.168.1.1", "pcap", False, False, [], True),
-        # # Testcase 4: Client IP
-        # ("conn", "http", "10.0.0.1", "pcap", False, False, ["10.0.0.1"], True),
-        # # Testcase 5: DoH server
-        # ("conn", "http", "1.1.1.1", "pcap", True, False, [], True),
+        # Testcase 1: Not a 'conn' flow type
+        ("dns", "dns", "8.8.8.8", "pcap", False, False, False, [], True),
+        # Testcase 2: DNS application protocol
+        ("conn", "dns", "8.8.8.8", "pcap", False, False, False, [], True),
+        # Testcase 3: Ignored IP
+        ("conn", "http", "192.168.1.1", "pcap", False, False, False, [], True),
+        # Testcase 4: Client IP
+        (
+            "conn",
+            "http",
+            "10.0.0.1",
+            "pcap",
+            False,
+            False,
+            False,
+            ["10.0.0.1"],
+            True,
+        ),
+        # Testcase 5: DoH server
+        ("conn", "http", "1.1.1.1", "pcap", True, False, False, [], True),
+        # Testcase 6: DHCP server
+        ("conn", "dhcp", "192.168.1.1", "pcap", True, False, True, [], True),
         # Testcase 7: Should not ignore
-        ("conn", "http", "93.184.216.34", "pcap", False, False, [], False),
+        (
+            "conn",
+            "http",
+            "93.184.216.34",
+            "pcap",
+            False,
+            False,
+            False,
+            [],
+            False,
+        ),
     ],
 )
 def test_should_ignore_conn_without_dns(
@@ -476,6 +498,7 @@ def test_should_ignore_conn_without_dns(
     input_type,
     is_doh_server,
     is_dns_server,
+    is_dhcp_server,
     client_ips,
     expected_result,
 ):
@@ -506,6 +529,7 @@ def test_should_ignore_conn_without_dns(
 
     conn.db.get_input_type.return_value = input_type
     conn.db.is_doh_server.return_value = is_doh_server
+    conn.db.is_dhcp_server.return_value = is_dhcp_server
     conn.dns_analyzer = Mock()
     conn.client_ips = client_ips
     mocker.patch(