1 files changed, 5 insertions, 6 deletions
diff --git a/bin/tests/system/statschannel/tests_xml.py b/bin/tests/system/statschannel/tests_xml.py
index 7f0b37e..5329615 100755
--- a/bin/tests/system/statschannel/tests_xml.py
+++ b/bin/tests/system/statschannel/tests_xml.py
@@ -16,9 +16,11 @@ import xml.etree.ElementTree as ET
 
 import pytest
 
-import generic
 import pytest_custom_markers
 
+pytest.register_assert_rewrite("generic")
+import generic
+
 pytestmark = pytest_custom_markers.have_libxml2
 requests = pytest.importorskip("requests")
 
@@ -128,8 +130,5 @@ def test_zone_with_many_keys_xml(statsport):
     )
 
 
-def test_traffic_xml(named_port, statsport):
-    generic_dnspython = pytest.importorskip("generic_dnspython")
-    generic_dnspython.test_traffic(
-        fetch_traffic_xml, statsip="10.53.0.2", statsport=statsport, port=named_port
-    )
+def test_traffic_xml(statsport):
+    generic.test_traffic(fetch_traffic_xml, statsip="10.53.0.2", statsport=statsport)