diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index a839926b8cf7..d204d2155d1c 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -520,8 +520,14 @@ in
   };
   firefoxpwa = runTest ./firefoxpwa.nix;
   firejail = runTest ./firejail.nix;
-  firewall = handleTest ./firewall.nix { nftables = false; };
-  firewall-nftables = handleTest ./firewall.nix { nftables = true; };
+  firewall = runTest {
+    imports = [ ./firewall.nix ];
+    _module.args.nftables = false;
+  };
+  firewall-nftables = runTest {
+    imports = [ ./firewall.nix ];
+    _module.args.nftables = true;
+  };
   fish = runTest ./fish.nix;
   firezone = runTest ./firezone/firezone.nix;
   flannel = handleTestOn [ "x86_64-linux" ] ./flannel.nix { };
diff --git a/nixos/tests/firewall.nix b/nixos/tests/firewall.nix
index 378f9b6a7c86..019508adb0b2 100644
--- a/nixos/tests/firewall.nix
+++ b/nixos/tests/firewall.nix
@@ -1,119 +1,117 @@
 # Test the firewall module.
 
-import ./make-test-python.nix (
-  { pkgs, nftables, ... }:
-  {
-    name = "firewall" + pkgs.lib.optionalString nftables "-nftables";
-    meta = with pkgs.lib.maintainers; {
-      maintainers = [
-        rvfg
-        garyguo
-      ];
-    };
+{ lib, nftables, ... }:
+{
+  name = "firewall" + lib.optionalString nftables "-nftables";
+  meta = with lib.maintainers; {
+    maintainers = [
+      rvfg
+      garyguo
+    ];
+  };
 
-    nodes = {
-      walled =
-        { ... }:
-        {
-          networking.firewall = {
-            enable = true;
-            logRefusedPackets = true;
-            # Syntax smoke test, not actually verified otherwise
-            allowedTCPPorts = [
-              25
-              993
-              8005
-            ];
+  nodes = {
+    walled =
+      { ... }:
+      {
+        networking.firewall = {
+          enable = true;
+          logRefusedPackets = true;
+          # Syntax smoke test, not actually verified otherwise
+          allowedTCPPorts = [
+            25
+            993
+            8005
+          ];
+          allowedTCPPortRanges = [
+            {
+              from = 980;
+              to = 1000;
+            }
+            {
+              from = 990;
+              to = 1010;
+            }
+            {
+              from = 8000;
+              to = 8010;
+            }
+          ];
+          interfaces.eth0 = {
+            allowedTCPPorts = [ 10003 ];
             allowedTCPPortRanges = [
               {
-                from = 980;
-                to = 1000;
-              }
-              {
-                from = 990;
-                to = 1010;
-              }
-              {
-                from = 8000;
-                to = 8010;
+                from = 10000;
+                to = 10005;
               }
             ];
-            interfaces.eth0 = {
-              allowedTCPPorts = [ 10003 ];
-              allowedTCPPortRanges = [
-                {
-                  from = 10000;
-                  to = 10005;
-                }
-              ];
-            };
-            interfaces.eth3 = {
-              allowedUDPPorts = [ 10003 ];
-              allowedUDPPortRanges = [
-                {
-                  from = 10000;
-                  to = 10005;
-                }
-              ];
-            };
           };
-          networking.nftables.enable = nftables;
-          services.httpd.enable = true;
-          services.httpd.adminAddr = "foo@example.org";
-
-          specialisation.different-config.configuration = {
-            networking.firewall.rejectPackets = true;
+          interfaces.eth3 = {
+            allowedUDPPorts = [ 10003 ];
+            allowedUDPPortRanges = [
+              {
+                from = 10000;
+                to = 10005;
+              }
+            ];
           };
         };
+        networking.nftables.enable = nftables;
+        services.httpd.enable = true;
+        services.httpd.adminAddr = "foo@example.org";
 
-      attacker =
-        { ... }:
-        {
-          services.httpd.enable = true;
-          services.httpd.adminAddr = "foo@example.org";
-          networking.firewall.enable = false;
+        specialisation.different-config.configuration = {
+          networking.firewall.rejectPackets = true;
         };
-    };
+      };
 
-    testScript =
-      { nodes, ... }:
-      let
-        unit = if nftables then "nftables" else "firewall";
-      in
-      ''
-        start_all()
+    attacker =
+      { ... }:
+      {
+        services.httpd.enable = true;
+        services.httpd.adminAddr = "foo@example.org";
+        networking.firewall.enable = false;
+      };
+  };
 
-        walled.wait_for_unit("${unit}")
-        walled.wait_for_unit("httpd")
-        attacker.wait_for_unit("network.target")
+  testScript =
+    { nodes, ... }:
+    let
+      unit = if nftables then "nftables" else "firewall";
+    in
+    ''
+      start_all()
 
-        # Local connections should still work.
-        walled.succeed("curl -v http://localhost/ >&2")
+      walled.wait_for_unit("${unit}")
+      walled.wait_for_unit("httpd")
+      attacker.wait_for_unit("network.target")
 
-        # Connections to the firewalled machine should fail, but ping should succeed.
-        attacker.fail("curl --fail --connect-timeout 2 http://walled/ >&2")
-        attacker.succeed("ping -c 1 walled >&2")
+      # Local connections should still work.
+      walled.succeed("curl -v http://localhost/ >&2")
 
-        # Outgoing connections/pings should still work.
-        walled.succeed("curl -v http://attacker/ >&2")
-        walled.succeed("ping -c 1 attacker >&2")
+      # Connections to the firewalled machine should fail, but ping should succeed.
+      attacker.fail("curl --fail --connect-timeout 2 http://walled/ >&2")
+      attacker.succeed("ping -c 1 walled >&2")
 
-        # Open tcp port 80 at runtime
-        walled.succeed("nixos-firewall-tool open tcp 80")
-        attacker.succeed("curl -v http://walled/ >&2")
+      # Outgoing connections/pings should still work.
+      walled.succeed("curl -v http://attacker/ >&2")
+      walled.succeed("ping -c 1 attacker >&2")
 
-        # Reset the firewall
-        walled.succeed("nixos-firewall-tool reset")
-        attacker.fail("curl --fail --connect-timeout 2 http://walled/ >&2")
+      # Open tcp port 80 at runtime
+      walled.succeed("nixos-firewall-tool open tcp 80")
+      attacker.succeed("curl -v http://walled/ >&2")
 
-        # If we stop the firewall, then connections should succeed.
-        walled.stop_job("${unit}")
-        attacker.succeed("curl -v http://walled/ >&2")
+      # Reset the firewall
+      walled.succeed("nixos-firewall-tool reset")
+      attacker.fail("curl --fail --connect-timeout 2 http://walled/ >&2")
 
-        # Check whether activation of a new configuration reloads the firewall.
-        walled.succeed(
-            "/run/booted-system/specialisation/different-config/bin/switch-to-configuration test 2>&1 | grep -qF ${unit}.service"
-        )
-      '';
-  }
-)
+      # If we stop the firewall, then connections should succeed.
+      walled.stop_job("${unit}")
+      attacker.succeed("curl -v http://walled/ >&2")
+
+      # Check whether activation of a new configuration reloads the firewall.
+      walled.succeed(
+          "/run/booted-system/specialisation/different-config/bin/switch-to-configuration test 2>&1 | grep -qF ${unit}.service"
+      )
+    '';
+}