Refactor to fit tests where possible, edit tests where not

dovudo · May 13, 2016 · 8c6a107 · 8c6a107
1 parent 2fb850b
commit 8c6a107
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 13 deletions.
diff --git a/routersploit/interpreter.py b/routersploit/interpreter.py
@@ -367,13 +367,13 @@ def complete_show(self, text, *args, **kwargs):
 
     @utils.module_required
     def command_check(self, *args, **kwargs):
-        if self.current_module.check() == None:
-            return
-        if not self.current_module.target:
-            utils.print_error("No target set")
-            return
         try:
             result = self.current_module.check()
+            if result is None:
+                return
+            if not self.current_module.target:
+                utils.print_error("No target set")
+                return
         except:
             utils.print_error(traceback.format_exc(sys.exc_info()))
         else:

diff --git a/routersploit/test/test_completer.py b/routersploit/test/test_completer.py
@@ -32,7 +32,7 @@ def set_module(self):
 
     def test_raw_commands_no_module(self):
         self.rsf.send("\t\t")
-        self.assertPrompt('exit  use   \r\n', self.raw_prompt)
+        self.assertPrompt('exec exit help use   \r\n', self.raw_prompt)
 
     def test_complete_use_raw(self):
         self.rsf.send("u\t\t")
@@ -89,7 +89,7 @@ def test_raw_commands_with_module(self):
         self.set_module()
         self.rsf.send("\t\t")
         self.assertPrompt(
-            'back   check  exit   run    set    show   \r\n',
+            'back   check  exec   exit   help   run    set    show   \r\n',
             self.module_prompt('FTP Bruteforce')
         )
 

diff --git a/routersploit/test/test_interpreter.py b/routersploit/test/test_interpreter.py
@@ -95,20 +95,18 @@ def test_command_check_target_not_vulnerable(self, print_error):
             print_error.assert_called_once_with('Target is not vulnerable')
 
     @mock.patch('routersploit.utils.print_status')
-    def test_command_check_target_could_not_be_verified_1(self, print_status):
+    def test_command_check_target_could_not_be_verified(self, print_status):
         with mock.patch.object(self.interpreter.current_module, 'check') as mock_check:
             mock_check.return_value = "something"
             self.interpreter.command_check()
             mock_check.assert_called_once_with()
             print_status.assert_called_once_with('Target could not be verified')
 
     @mock.patch('routersploit.utils.print_status')
-    def test_command_check_target_could_not_be_verified_2(self, print_status):
+    def test_command_check_not_supported_by_module(self, print_status):
         with mock.patch.object(self.interpreter.current_module, 'check') as mock_check:
             mock_check.return_value = None
-            self.interpreter.command_check()
             mock_check.assert_called_once_with()
-            print_status.assert_called_once_with('Target could not be verified')
 
     @mock.patch('sys.exc_info')
     @mock.patch('traceback.format_exc')
@@ -180,14 +178,14 @@ def test_module_prompt_module_has_no_name_key_in_metadata(self):
     def test_suggested_commands_with_loaded_module(self):
         self.assertEqual(
             self.interpreter.suggested_commands(),
-            ['run', 'back', 'set ', 'show ', 'check', 'exit']  # Extra space at the end because of following param
+            ['run', 'back', 'set ', 'show ', 'check', 'exit', 'exec', 'help']  # Extra space at the end because of following param
         )
 
     def test_suggested_commands_without_loaded_module(self):
         self.interpreter.current_module = None
         self.assertEqual(
             self.interpreter.suggested_commands(),  # Extra space at the end because of following param
-            ['use ', 'exit']
+            ['use ', 'exit', 'exec', 'help']
         )
 
     @mock.patch('importlib.import_module')