@@ -8,72 +8,81 @@ class MainTest(BaseTestCase):
88 @mock .patch ('pyt.__main__.discover_files' )
99 @mock .patch ('pyt.__main__.parse_args' )
1010 @mock .patch ('pyt.__main__.find_vulnerabilities' )
11- @mock .patch ('pyt.__main__ .text' )
11+ @mock .patch ('pyt.formatters .text' )
1212 def test_text_output (self , mock_text , mock_find_vulnerabilities , mock_parse_args , mock_discover_files ):
1313 mock_find_vulnerabilities .return_value = 'stuff'
1414 example_file = 'examples/vulnerable_code/inter_command_injection.py'
1515 output_file = 'mocked_outfile'
1616
17+ import pyt .formatters .text
1718 mock_discover_files .return_value = [example_file ]
1819 mock_parse_args .return_value = mock .Mock (
1920 project_root = None ,
2021 baseline = None ,
21- json = None ,
22- output_file = output_file
22+ formatter = pyt .formatters .text ,
23+ output_file = output_file ,
24+ only_unsanitised = False ,
2325 )
2426 with self .assertRaises (SystemExit ):
2527 main (['parse_args is mocked' ])
2628 assert mock_text .report .call_count == 1
2729 mock_text .report .assert_called_with (
2830 mock_find_vulnerabilities .return_value ,
29- mock_parse_args .return_value .output_file
31+ mock_parse_args .return_value .output_file ,
32+ True ,
3033 )
3134
3235 @mock .patch ('pyt.__main__.discover_files' )
3336 @mock .patch ('pyt.__main__.parse_args' )
3437 @mock .patch ('pyt.__main__.find_vulnerabilities' )
35- @mock .patch ('pyt.__main__ .text' )
38+ @mock .patch ('pyt.formatters .text' )
3639 def test_no_vulns_found (self , mock_text , mock_find_vulnerabilities , mock_parse_args , mock_discover_files ):
3740 mock_find_vulnerabilities .return_value = []
3841 example_file = 'examples/vulnerable_code/inter_command_injection.py'
3942 output_file = 'mocked_outfile'
4043
44+ import pyt .formatters .text
4145 mock_discover_files .return_value = [example_file ]
4246 mock_parse_args .return_value = mock .Mock (
4347 project_root = None ,
4448 baseline = None ,
45- json = None ,
46- output_file = output_file
49+ formatter = pyt .formatters .text ,
50+ output_file = output_file ,
51+ only_unsanitised = True ,
4752 )
4853 main (['parse_args is mocked' ]) # No SystemExit
4954 assert mock_text .report .call_count == 1
5055 mock_text .report .assert_called_with (
5156 mock_find_vulnerabilities .return_value ,
52- mock_parse_args .return_value .output_file
57+ mock_parse_args .return_value .output_file ,
58+ False ,
5359 )
5460
5561 @mock .patch ('pyt.__main__.discover_files' )
5662 @mock .patch ('pyt.__main__.parse_args' )
5763 @mock .patch ('pyt.__main__.find_vulnerabilities' )
58- @mock .patch ('pyt.__main__ .json' )
64+ @mock .patch ('pyt.formatters .json' )
5965 def test_json_output (self , mock_json , mock_find_vulnerabilities , mock_parse_args , mock_discover_files ):
6066 mock_find_vulnerabilities .return_value = 'stuff'
6167 example_file = 'examples/vulnerable_code/inter_command_injection.py'
6268 output_file = 'mocked_outfile'
6369
70+ import pyt .formatters .json
6471 mock_discover_files .return_value = [example_file ]
6572 mock_parse_args .return_value = mock .Mock (
6673 project_root = None ,
6774 baseline = None ,
68- json = True ,
69- output_file = output_file
75+ formatter = pyt .formatters .json ,
76+ output_file = output_file ,
77+ only_unsanitised = False ,
7078 )
7179 with self .assertRaises (SystemExit ):
7280 main (['parse_args is mocked' ])
7381 assert mock_json .report .call_count == 1
7482 mock_json .report .assert_called_with (
7583 mock_find_vulnerabilities .return_value ,
76- mock_parse_args .return_value .output_file
84+ mock_parse_args .return_value .output_file ,
85+ True ,
7786 )
7887
7988
0 commit comments