From 9cc6d636843179d836222e1849994375b199bb52 Mon Sep 17 00:00:00 2001
From: Brian May <brian@linuxpenguins.xyz>
Date: Tue, 17 Nov 2015 12:46:35 +1100
Subject: [PATCH] Fix firewall tests.

---
 sshuttle/tests/test_firewall.py | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/sshuttle/tests/test_firewall.py b/sshuttle/tests/test_firewall.py
index dca1e5d..540a5bf 100644
--- a/sshuttle/tests/test_firewall.py
+++ b/sshuttle/tests/test_firewall.py
@@ -24,6 +24,11 @@ GO 1
     return stdin, stdout
 
 
+@patch('sshuttle.firewall.HOSTSFILE', new='tmp/hosts')
+@patch('sshuttle.firewall.hostmap', new={
+    'myhost': '1.2.3.4',
+    'myotherhost': '1.2.3.5',
+})
 def test_rewrite_etc_hosts():
     if not os.path.isdir("tmp"):
         os.mkdir("tmp")
@@ -33,11 +38,6 @@ def test_rewrite_etc_hosts():
 
     shutil.copyfile("tmp/hosts.orig", "tmp/hosts")
 
-    sshuttle.firewall.HOSTSFILE = "tmp/hosts"
-    sshuttle.firewall.hostmap = {
-        'myhost': '1.2.3.4',
-        'myotherhost': '1.2.3.5',
-    }
     sshuttle.firewall.rewrite_etc_hosts(10)
     with open("tmp/hosts") as f:
         line = f.readline()
@@ -61,14 +61,26 @@ def test_rewrite_etc_hosts():
     assert filecmp.cmp("tmp/hosts.orig", "tmp/hosts", shallow=False) is True
 
 
+@patch('sshuttle.firewall.HOSTSFILE', new='tmp/hosts')
 @patch('sshuttle.firewall.setup_daemon')
 @patch('sshuttle.firewall.get_method')
 def test_main(mock_get_method, mock_setup_daemon):
     stdin, stdout = setup_daemon()
     mock_setup_daemon.return_value = stdin, stdout
 
+    if not os.path.isdir("tmp"):
+        os.mkdir("tmp")
+
     sshuttle.firewall.main("test", False)
 
+    with open("tmp/hosts") as f:
+        line = f.readline()
+        s = line.split()
+        assert s == ['1.2.3.3', 'existing']
+
+        line = f.readline()
+        assert line == ""
+
     stdout.mock_calls == [
         call.write('READY test\n'),
         call.flush(),