|
1 | 1 | from mock import Mock, patch, call
|
2 |
| -from contextlib import nested |
3 | 2 | import io
|
4 | 3 | import os
|
5 | 4 | import os.path
|
@@ -63,39 +62,37 @@ def test_rewrite_etc_hosts():
|
63 | 62 | assert filecmp.cmp("tmp/hosts.orig", "tmp/hosts", shallow=False) is True
|
64 | 63 |
|
65 | 64 |
|
66 |
| -def test_main(): |
67 |
| - with nested( |
68 |
| - patch('sshuttle.firewall.setup_daemon'), |
69 |
| - patch('sshuttle.firewall.get_method') |
70 |
| - ) as (mock_setup_daemon, mock_get_method): |
71 |
| - stdin, stdout = setup_daemon() |
72 |
| - mock_setup_daemon.return_value = stdin, stdout |
73 |
| - |
74 |
| - sshuttle.firewall.main("test", False) |
75 |
| - |
76 |
| - stdout.mock_calls == [ |
77 |
| - call.write('READY test\n'), |
78 |
| - call.flush(), |
79 |
| - call.write('STARTED\n'), |
80 |
| - call.flush() |
81 |
| - ] |
82 |
| - mock_setup_daemon.mock_calls == [call()] |
83 |
| - mock_get_method.mock_calls == [ |
84 |
| - call('test'), |
85 |
| - call().setup_firewall( |
86 |
| - 1024, 1026, |
87 |
| - [(10, u'2404:6800:4004:80c::33')], |
88 |
| - 10, |
89 |
| - [(10, 64, False, u'2404:6800:4004:80c::'), |
90 |
| - (10, 128, True, u'2404:6800:4004:80c::101f')], |
91 |
| - True), |
92 |
| - call().setup_firewall( |
93 |
| - 1025, 1027, |
94 |
| - [(2, u'1.2.3.33')], |
95 |
| - 2, |
96 |
| - [(2, 24, False, u'1.2.3.0'), (2, 32, True, u'1.2.3.66')], |
97 |
| - True), |
98 |
| - call().setup_firewall()(), |
99 |
| - call().setup_firewall(1024, 0, [], 10, [], True), |
100 |
| - call().setup_firewall(1025, 0, [], 2, [], True), |
101 |
| - ] |
| 65 | +@patch('sshuttle.firewall.setup_daemon') |
| 66 | +@patch('sshuttle.firewall.get_method') |
| 67 | +def test_main(mock_get_method, mock_setup_daemon): |
| 68 | + stdin, stdout = setup_daemon() |
| 69 | + mock_setup_daemon.return_value = stdin, stdout |
| 70 | + |
| 71 | + sshuttle.firewall.main("test", False) |
| 72 | + |
| 73 | + stdout.mock_calls == [ |
| 74 | + call.write('READY test\n'), |
| 75 | + call.flush(), |
| 76 | + call.write('STARTED\n'), |
| 77 | + call.flush() |
| 78 | + ] |
| 79 | + mock_setup_daemon.mock_calls == [call()] |
| 80 | + mock_get_method.mock_calls == [ |
| 81 | + call('test'), |
| 82 | + call().setup_firewall( |
| 83 | + 1024, 1026, |
| 84 | + [(10, u'2404:6800:4004:80c::33')], |
| 85 | + 10, |
| 86 | + [(10, 64, False, u'2404:6800:4004:80c::'), |
| 87 | + (10, 128, True, u'2404:6800:4004:80c::101f')], |
| 88 | + True), |
| 89 | + call().setup_firewall( |
| 90 | + 1025, 1027, |
| 91 | + [(2, u'1.2.3.33')], |
| 92 | + 2, |
| 93 | + [(2, 24, False, u'1.2.3.0'), (2, 32, True, u'1.2.3.66')], |
| 94 | + True), |
| 95 | + call().setup_firewall()(), |
| 96 | + call().setup_firewall(1024, 0, [], 10, [], True), |
| 97 | + call().setup_firewall(1025, 0, [], 2, [], True), |
| 98 | + ] |
0 commit comments