@@ -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