[PRED-2176] Fix tests

datarobot · Mar 22, 2019 · f2dde30 · f2dde30
1 parent 2e0d128
commit f2dde30
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/tests/test_api_response_handlers.py b/tests/test_api_response_handlers.py
@@ -164,21 +164,23 @@ def test_unpack_data(self):
 
         ({'pred_name': None,
           'pred_decision_name': 'label',
+          'pred_threshold_name': None,
           'keep_cols': None,
           'skip_row_id': False,
           'fast_mode': False,
           'delimiter': ','},
          ['row_id', 'readmitted', 'label'],
          [[0, 1, 1], [1, 1, 1]]),
 
-        ({'pred_name': 'threshold',
+        ({'pred_name': None,
           'pred_decision_name': 'label',
+          'pred_threshold_name': 'threshold',
           'keep_cols': None,
           'skip_row_id': False,
           'fast_mode': False,
           'delimiter': ','},
-         ['row_id', 'readmitted', 'label', 'threshold'],
-         [[0, 1, 1, 0.5], [1, 1, 1, 0.5]]),
+         ['row_id', 'readmitted', 'threshold', 'label'],
+         [[0, 1, 0.5, 1], [1, 1, 0.5, 1]]),
     ))
     def test_format_data(self, parsed_pred_api_v10_predictions, batch,
                          opts, expected_fields, expected_values):

diff --git a/tests/test_main.py b/tests/test_main.py
@@ -495,6 +495,7 @@ def test_pred_threshold(monkeypatch):
             delimiter=None,
             dataset='tests/fixtures/temperatura_predict.csv',
             pred_name=None,
+            pred_decision_name=None,
             pred_threshold_name='threshold',
             timeout=None,
             ui=mock.ANY,
@@ -548,6 +549,7 @@ def test_pred_decision(monkeypatch):
             dataset='tests/fixtures/temperatura_predict.csv',
             pred_name=None,
             pred_decision_name='label',
+            pred_threshold_name=None,
             timeout=None,
             ui=mock.ANY,
             fast_mode=False,