Fix flake8 issues

NannyML · Apr 29, 2024 · b6d8bdf · b6d8bdf
1 parent 4eb8cb8
commit b6d8bdf
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 22 deletions.
diff --git a/tests/performance_calculation/test_performance_calculator.py b/tests/performance_calculation/test_performance_calculator.py
@@ -466,13 +466,15 @@ def test_binary_classification_result_plots_raise_no_exceptions(calc_args, plot_
 
 def test_binary_classification_calculate_without_prediction_column():
     reference, analysis, analysis_targets = load_synthetic_binary_classification_dataset()
-    calc = PerformanceCalculator(
-        y_true='work_home_actual',
-        y_pred_proba='y_pred_proba',
-        problem_type=ProblemType.CLASSIFICATION_BINARY,
-        metrics=['roc_auc', 'average_precision'],
-        timestamp_column_name='timestamp',
-        chunk_period='M'
-    ).fit(reference)
-    res = calc.calculate(analysis.merge(analysis_targets, on='id'))
-
+    try:
+        calc = PerformanceCalculator(
+            y_true='work_home_actual',
+            y_pred_proba='y_pred_proba',
+            problem_type=ProblemType.CLASSIFICATION_BINARY,
+            metrics=['roc_auc', 'average_precision'],
+            timestamp_column_name='timestamp',
+            chunk_period='M',
+        ).fit(reference)
+        _ = calc.calculate(analysis.merge(analysis_targets, on='id'))
+    except Exception as exc:
+        pytest.fail(f"an unexpected exception occurred: {exc}")
diff --git a/tests/performance_estimation/CBPE/test_cbpe.py b/tests/performance_estimation/CBPE/test_cbpe.py
@@ -706,15 +706,18 @@ def test_cbpe_with_default_thresholds():
 
 def test_cbpe_without_predictions():
     ref_df, ana_df, _ = load_synthetic_binary_classification_dataset()
-    cbpe = CBPE(
-        y_pred_proba='y_pred_proba',
-        y_true='work_home_actual',
-        problem_type='classification_binary',
-        metrics=[
-            'roc_auc',
-            'average_precision',
-        ],
-        timestamp_column_name='timestamp',
-        chunk_period='M',
-    ).fit(ref_df)
-    result = cbpe.estimate(ana_df)
+    try:
+        cbpe = CBPE(
+            y_pred_proba='y_pred_proba',
+            y_true='work_home_actual',
+            problem_type='classification_binary',
+            metrics=[
+                'roc_auc',
+                'average_precision',
+            ],
+            timestamp_column_name='timestamp',
+            chunk_period='M',
+        ).fit(ref_df)
+        _ = cbpe.estimate(ana_df)
+    except Exception as exc:
+        pytest.fail(f'unexpected exception: {exc}')