Fix tests

4 files changed, 27 insertions, 11 deletions

diff --git a/tests/client/test_firewall.py b/tests/client/test_firewall.py
index 7404328..6a3b4db 100644
--- a/tests/client/test_firewall.py
+++ b/tests/client/test_firewall.py
@@ -145,6 +145,7 @@ def test_main(mock_get_method, mock_setup_daemon, mock_rewrite_etc_hosts):
              (AF_INET6, 128, True, u'2404:6800:4004:80c::101f', 80, 80)],
             True,
             None,
+            None,
             '0x01'),
         call().setup_firewall(
             1025, 1027,
@@ -154,7 +155,8 @@ def test_main(mock_get_method, mock_setup_daemon, mock_rewrite_etc_hosts):
              (AF_INET, 32, True, u'1.2.3.66', 8080, 8080)],
             True,
             None,
+            None,
             '0x01'),
-        call().restore_firewall(1024, AF_INET6, True, None),
-        call().restore_firewall(1025, AF_INET, True, None),
+        call().restore_firewall(1024, AF_INET6, True, None, None),
+        call().restore_firewall(1025, AF_INET, True, None, None),
     ]
diff --git a/tests/client/test_methods_nat.py b/tests/client/test_methods_nat.py
index 61b74b5..218b629 100644
--- a/tests/client/test_methods_nat.py
+++ b/tests/client/test_methods_nat.py
@@ -101,6 +101,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
          (AF_INET6, 128, True, u'2404:6800:4004:80c::101f', 80, 80)],
         False,
         None,
+        None,
         '0x01')
 
     assert mock_ipt_chain_exists.mock_calls == [
@@ -142,6 +143,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
                 (AF_INET, 32, True, u'1.2.3.66', 8080, 8080)],
             True,
             None,
+            None,
             '0x01')
     assert str(excinfo.value) == 'UDP not supported by nat method_name'
     assert mock_ipt_chain_exists.mock_calls == []
@@ -155,6 +157,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
             (AF_INET, 32, True, u'1.2.3.66', 8080, 8080)],
         False,
         None,
+        None,
         '0x01')
     assert mock_ipt_chain_exists.mock_calls == [
         call(AF_INET, 'nat', 'sshuttle-1025')
@@ -182,7 +185,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
     mock_ipt_chain_exists.reset_mock()
     mock_ipt.reset_mock()
 
-    method.restore_firewall(1025, AF_INET, False, None)
+    method.restore_firewall(1025, AF_INET, False, None, None)
     assert mock_ipt_chain_exists.mock_calls == [
         call(AF_INET, 'nat', 'sshuttle-1025')
     ]
@@ -197,7 +200,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
     mock_ipt_chain_exists.reset_mock()
     mock_ipt.reset_mock()
 
-    method.restore_firewall(1025, AF_INET6, False, None)
+    method.restore_firewall(1025, AF_INET6, False, None,)
     assert mock_ipt_chain_exists.mock_calls == [
         call(AF_INET6, 'nat', 'sshuttle-1025')
     ]
diff --git a/tests/client/test_methods_pf.py b/tests/client/test_methods_pf.py
index ea7416a..46de24c 100644
--- a/tests/client/test_methods_pf.py
+++ b/tests/client/test_methods_pf.py
@@ -187,6 +187,7 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl):
             (AF_INET6, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)],
         False,
         None,
+        None,
         '0x01')
     assert mock_ioctl.mock_calls == [
         call(mock_pf_get_dev(), 0xC4704433, ANY),
@@ -227,6 +228,7 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl):
                 (AF_INET, 32, True, u'1.2.3.66', 80, 80)],
             True,
             None,
+            None,
             '0x01')
     assert str(excinfo.value) == 'UDP not supported by pf method_name'
     assert mock_pf_get_dev.mock_calls == []
@@ -241,6 +243,7 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl):
             (AF_INET, 32, True, u'1.2.3.66', 80, 80)],
         False,
         None,
+        None,
         '0x01')
     assert mock_ioctl.mock_calls == [
         call(mock_pf_get_dev(), 0xC4704433, ANY),
@@ -270,7 +273,7 @@ def test_setup_firewall_darwin(mock_pf_get_dev, mock_ioctl, mock_pfctl):
     mock_ioctl.reset_mock()
     mock_pfctl.reset_mock()
 
-    method.restore_firewall(1025, AF_INET, False, None)
+    method.restore_firewall(1025, AF_INET, False, None, None)
     assert mock_ioctl.mock_calls == []
     assert mock_pfctl.mock_calls == [
         call('-a sshuttle-1025 -F all'),
@@ -302,6 +305,7 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl,
             (AF_INET6, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)],
         False,
         None,
+        None,
         '0x01')
 
     assert mock_pfctl.mock_calls == [
@@ -335,6 +339,7 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl,
                 (AF_INET, 32, True, u'1.2.3.66', 80, 80)],
             True,
             None,
+            None,
             '0x01')
     assert str(excinfo.value) == 'UDP not supported by pf method_name'
     assert mock_pf_get_dev.mock_calls == []
@@ -349,6 +354,7 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl,
             (AF_INET, 32, True, u'1.2.3.66', 80, 80)],
         False,
         None,
+        None,
         '0x01')
     assert mock_ioctl.mock_calls == [
         call(mock_pf_get_dev(), 0xC4704433, ANY),
@@ -376,8 +382,8 @@ def test_setup_firewall_freebsd(mock_pf_get_dev, mock_ioctl, mock_pfctl,
     mock_ioctl.reset_mock()
     mock_pfctl.reset_mock()
 
-    method.restore_firewall(1025, AF_INET, False, None)
-    method.restore_firewall(1024, AF_INET6, False, None)
+    method.restore_firewall(1025, AF_INET, False, None, None)
+    method.restore_firewall(1024, AF_INET6, False, None, None)
     assert mock_ioctl.mock_calls == []
     assert mock_pfctl.mock_calls == [
         call('-a sshuttle-1025 -F all'),
@@ -408,6 +414,7 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl):
             (AF_INET6, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)],
         False,
         None,
+        None,
         '0x01')
 
     assert mock_ioctl.mock_calls == [
@@ -445,6 +452,7 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl):
                 (AF_INET, 32, True, u'1.2.3.66', 80, 80)],
             True,
             None,
+            None,
             '0x01')
     assert str(excinfo.value) == 'UDP not supported by pf method_name'
     assert mock_pf_get_dev.mock_calls == []
@@ -459,6 +467,7 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl):
             (AF_INET, 32, True, u'1.2.3.66', 80, 80)],
         False,
         None,
+        None,
         '0x01')
     assert mock_ioctl.mock_calls == [
         call(mock_pf_get_dev(), 0xcd60441a, ANY),
@@ -484,8 +493,8 @@ def test_setup_firewall_openbsd(mock_pf_get_dev, mock_ioctl, mock_pfctl):
     mock_ioctl.reset_mock()
     mock_pfctl.reset_mock()
 
-    method.restore_firewall(1025, AF_INET, False, None)
-    method.restore_firewall(1024, AF_INET6, False, None)
+    method.restore_firewall(1025, AF_INET, False, None, None)
+    method.restore_firewall(1024, AF_INET6, False, None, None)
     assert mock_ioctl.mock_calls == []
     assert mock_pfctl.mock_calls == [
         call('-a sshuttle-1025 -F all'),
diff --git a/tests/client/test_methods_tproxy.py b/tests/client/test_methods_tproxy.py
index d7bc33f..b1e72ff 100644
--- a/tests/client/test_methods_tproxy.py
+++ b/tests/client/test_methods_tproxy.py
@@ -98,6 +98,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
             (AF_INET6, 128, True, u'2404:6800:4004:80c::101f', 8080, 8080)],
         True,
         None,
+        None,
         '0x01')
     assert mock_ipt_chain_exists.mock_calls == [
         call(AF_INET6, 'mangle', 'sshuttle-m-1024'),
@@ -172,7 +173,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
     mock_ipt_chain_exists.reset_mock()
     mock_ipt.reset_mock()
 
-    method.restore_firewall(1025, AF_INET6, True, None)
+    method.restore_firewall(1025, AF_INET6, True, None, None)
     assert mock_ipt_chain_exists.mock_calls == [
         call(AF_INET6, 'mangle', 'sshuttle-m-1025'),
         call(AF_INET6, 'mangle', 'sshuttle-t-1025'),
@@ -201,6 +202,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
             (AF_INET, 32, True, u'1.2.3.66', 80, 80)],
         True,
         None,
+        None,
         '0x01')
     assert mock_ipt_chain_exists.mock_calls == [
         call(AF_INET, 'mangle', 'sshuttle-m-1025'),
@@ -270,7 +272,7 @@ def test_setup_firewall(mock_ipt_chain_exists, mock_ipt):
     mock_ipt_chain_exists.reset_mock()
     mock_ipt.reset_mock()
 
-    method.restore_firewall(1025, AF_INET, True, None)
+    method.restore_firewall(1025, AF_INET, True, None, None)
     assert mock_ipt_chain_exists.mock_calls == [
         call(AF_INET, 'mangle', 'sshuttle-m-1025'),
         call(AF_INET, 'mangle', 'sshuttle-t-1025'),