Improve coverage

.coveragerc

@@ -0,0 +1,7 @@
+[run]
+branch = True
+source = datarobot_batch_scoring, tests
+omit = site-packages
+
+[html]
+directory = htmlcov

datarobot_batch_scoring/main.py

@@ -44,6 +44,7 @@ def main(argv=sys.argv[1:]):
     parser = argparse.ArgumentParser(
         description=DESCRIPTION, epilog=EPILOG,
         formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.set_defaults(prompt=None)
     parser.add_argument('--verbose', '-v', action="store_true",
                         help='Provides status updates while '
                         'the script is running.')
@@ -147,7 +148,6 @@ def main(argv=sys.argv[1:]):
     csv_gr.add_argument('--pred_name', type=str,
                         nargs='?')
     misc_gr = parser.add_argument_group('Miscellaneous')
-    misc_gr.set_defaults(prompt=None)
     misc_gr.add_argument('-y', '--yes', dest='prompt', action='store_true',
                          help="Always answer 'yes' for user prompts")
     misc_gr.add_argument('-n', '--no', dest='prompt', action='store_false',
@@ -163,9 +163,9 @@ def main(argv=sys.argv[1:]):
                        if v is not None}
     parsed_args.update(pre_parsed_args)
     loglevel = logging.DEBUG if parsed_args['verbose'] else logging.INFO
-    ui = UI(parsed_args['prompt'], loglevel)
+    ui = UI(parsed_args.get('prompt'), loglevel)
     printed_args = copy.copy(parsed_args)
-    printed_args.pop('password')
+    printed_args.pop('password', None)
     ui.debug(printed_args)
     ui.info('platform: {} {}'.format(sys.platform, sys.version))
 
@@ -185,7 +185,6 @@ def main(argv=sys.argv[1:]):
     resume = parsed_args['resume']
     out_file = parsed_args['out']
     datarobot_key = parsed_args.get('datarobot_key')
-    pwd = parsed_args['password']
     timeout = int(parsed_args['timeout'])
 
     if 'user' not in parsed_args:
@@ -200,11 +199,11 @@ def main(argv=sys.argv[1:]):
         verify_objectid(pid)
         verify_objectid(lid)
     except ValueError as e:
-        ui.fatal('{}'.format(e))
+        ui.fatal(str(e))
 
     api_token = parsed_args.get('api_token')
     create_api_token = parsed_args.get('create_api_token')
-    pwd = parsed_args['password']
+    pwd = parsed_args.get('password')
     pred_name = parsed_args.get('pred_name')
 
     api_version = parsed_args['api_version']
@@ -244,4 +243,4 @@ def main(argv=sys.argv[1:]):
 
 
 if __name__ == '__main__':
-    main()
+    main()  # pragma: no cover

tests/test_main.py

@@ -0,0 +1,160 @@
+import mock
+import pytest
+from datarobot_batch_scoring.main import main, UI
+
+
+def test_without_passed_user_and_passwd(monkeypatch):
+    main_args = ['--host',
+                 'http://localhost:53646/api',
+                 '56dd9570018e213242dfa93c',
+                 '56dd9570018e213242dfa93d',
+                 'tests/fixtures/temperatura_predict.csv',
+                 '--n_samples',
+                 '10',
+                 '--n_concurrent', '1', '--no']
+
+    monkeypatch.setattr('datarobot_batch_scoring.main.UI', mock.Mock(spec=UI))
+
+    with mock.patch(
+            'datarobot_batch_scoring.main'
+            '.run_batch_predictions_v1') as mock_method:
+        main(argv=main_args)
+        mock_method.assert_called_once_with(
+            base_url='http://localhost:53646/api/v1/',
+            base_headers={},
+            user=mock.ANY,
+            pwd=None,
+            api_token=None,
+            create_api_token=False,
+            pid='56dd9570018e213242dfa93c',
+            lid='56dd9570018e213242dfa93d',
+            n_retry=3,
+            concurrent=1,
+            resume=False,
+            n_samples=10,
+            out_file='out.csv',
+            keep_cols=None,
+            delimiter=None,
+            dataset='tests/fixtures/temperatura_predict.csv',
+            pred_name=None,
+            timeout=30,
+            ui=mock.ANY
+        )
+
+
+def test_keep_cols(monkeypatch):
+    main_args = ['--host',
+                 'http://localhost:53646/api',
+                 '56dd9570018e213242dfa93c',
+                 '56dd9570018e213242dfa93d',
+                 'tests/fixtures/temperatura_predict.csv',
+                 '--keep_cols', 'a, b, c']
+
+    monkeypatch.setattr('datarobot_batch_scoring.main.UI', mock.Mock(spec=UI))
+
+    with mock.patch(
+            'datarobot_batch_scoring.main'
+            '.run_batch_predictions_v1') as mock_method:
+        main(argv=main_args)
+        mock_method.assert_called_once_with(
+            base_url='http://localhost:53646/api/v1/',
+            base_headers={},
+            user=mock.ANY,
+            pwd=None,
+            api_token=None,
+            create_api_token=False,
+            pid='56dd9570018e213242dfa93c',
+            lid='56dd9570018e213242dfa93d',
+            n_retry=3,
+            concurrent=4,
+            resume=False,
+            n_samples=1000,
+            out_file='out.csv',
+            keep_cols=['a', 'b', 'c'],
+            delimiter=None,
+            dataset='tests/fixtures/temperatura_predict.csv',
+            pred_name=None,
+            timeout=30,
+            ui=mock.ANY
+        )
+
+
+def test_input_dataset_doesnt_exist(monkeypatch):
+    main_args = ['--host',
+                 'http://localhost:53646/api',
+                 '56dd9570018e213242dfa93c',
+                 '56dd9570018e213242dfa93d',
+                 'file-not-exists.csv']
+
+    ui_class = mock.Mock(spec=UI)
+    ui = ui_class.return_value
+    ui.fatal.side_effect = SystemExit
+    monkeypatch.setattr('datarobot_batch_scoring.main.UI', ui_class)
+
+    with mock.patch(
+            'datarobot_batch_scoring.main'
+            '.run_batch_predictions_v1') as mock_method:
+        with pytest.raises(SystemExit):
+            main(argv=main_args)
+        assert not mock_method.called
+    ui.fatal.assert_called_with('file file-not-exists.csv does not exist.')
+
+
+def test_bad_objectid(monkeypatch):
+    main_args = ['--host',
+                 'http://localhost:53646/api',
+                 '56dd9570018e213242dfa93caa',
+                 '56dd9570018e213242dfa93d',
+                 'tests/fixtures/temperatura_predict.csv']
+
+    ui_class = mock.Mock(spec=UI)
+    ui = ui_class.return_value
+    ui.fatal.side_effect = SystemExit
+    monkeypatch.setattr('datarobot_batch_scoring.main.UI', ui_class)
+    monkeypatch.setattr('datarobot_batch_scoring.main.verify_objectid',
+                        mock.Mock(side_effect=ValueError('bad objectid')))
+
+    with mock.patch(
+            'datarobot_batch_scoring.main'
+            '.run_batch_predictions_v1') as mock_method:
+        with pytest.raises(SystemExit):
+            main(argv=main_args)
+        assert not mock_method.called
+    ui.fatal.assert_called_with('bad objectid')
+
+
+def test_datarobot_key(monkeypatch):
+    main_args = ['--host',
+                 'http://localhost:53646/api',
+                 '56dd9570018e213242dfa93c',
+                 '56dd9570018e213242dfa93d',
+                 'tests/fixtures/temperatura_predict.csv',
+                 '--datarobot_key', 'the_key']
+
+    monkeypatch.setattr('datarobot_batch_scoring.main.UI', mock.Mock(spec=UI))
+
+    with mock.patch(
+            'datarobot_batch_scoring.main'
+            '.run_batch_predictions_v1') as mock_method:
+        main(argv=main_args)
+        mock_method.assert_called_once_with(
+            base_url='http://localhost:53646/api/v1/',
+            base_headers={'datarobot-key': 'the_key'},
+            user=mock.ANY,
+            pwd=None,
+            api_token=None,
+            create_api_token=False,
+            pid='56dd9570018e213242dfa93c',
+            lid='56dd9570018e213242dfa93d',
+            n_retry=3,
+            concurrent=4,
+            resume=False,
+            n_samples=1000,
+            out_file='out.csv',
+            keep_cols=None,
+            delimiter=None,
+            dataset='tests/fixtures/temperatura_predict.csv',
+            pred_name=None,
+            timeout=30,
+            ui=mock.ANY
+        )