update unit tests

Author: ahosler

tests/operators/forecast/test_errors.py

@@ -948,6 +948,38 @@ def test_prophet_floor_cap(operator_setup, model):
     ), "`max` not obeyed in prophet"
 
 
+def _check_results_obj(results):
+    assert not results.get_forecast().empty
+    assert not results.get_metrics().empty
+    assert not results.get_global_explanations().empty
+    assert not results.get_local_explanations().empty
+
+
+def _check_no_skippable_files(yaml_i, check_report=True):
+    files = os.listdir(yaml_i["spec"]["output_directory"]["url"])
+
+    if "errors.json" in files:
+        with open(
+            os.path.join(yaml_i["spec"]["output_directory"]["url"], "errors.json")
+        ) as f:
+            assert False, f"Failed due to errors.json being created: {f.read()}"
+    if check_report:
+        assert "report.html" in files, "Failed to generate report"
+
+    assert (
+        "forecast.csv" not in files
+    ), "Generated forecast file, but `generate_forecast_file` was set False"
+    assert (
+        "metrics.csv" not in files
+    ), "Generated metrics file, but `generate_metrics_file` was set False"
+    assert (
+        "local_explanations.csv" not in files
+    ), "Generated metrics file, but `generate_explanation_files` was set False"
+    assert (
+        "global_explanations.csv" not in files
+    ), "Generated metrics file, but `generate_explanation_files` was set False"
+
+
 @pytest.mark.parametrize("model", ["prophet"])
 def test_generate_files(operator_setup, model):
     yaml_i = TEMPLATE_YAML.copy()
@@ -970,35 +1002,15 @@ def test_generate_files(operator_setup, model):
     yaml_i["spec"]["additional_data"]["data"] = df_add
     operator_config = ForecastOperatorConfig.from_dict(yaml_i)
     results = operate(operator_config)
-    files = os.listdir(yaml_i["spec"]["output_directory"]["url"])
-    if "errors.json" in files:
-        with open(
-            os.path.join(yaml_i["spec"]["output_directory"]["url"], "errors.json")
-        ) as f:
-            assert False, f"Failed due to errors.json being created: {f.read()}"
-    assert "report.html" in files, "Failed to generate report"
-    assert (
-        "forecast.csv" not in files
-    ), "Generated forecast file, but `generate_forecast_file` was set False"
-    assert (
-        "metrics.csv" not in files
-    ), "Generated metrics file, but `generate_metrics_file` was set False"
-    assert (
-        "local_explanations.csv" not in files
-    ), "Generated metrics file, but `generate_explanation_files` was set False"
-    assert (
-        "global_explanations.csv" not in files
-    ), "Generated metrics file, but `generate_explanation_files` was set False"
-    assert not results.get_forecast().empty
-    assert not results.get_metrics().empty
-    assert not results.get_global_explanations().empty
-    assert not results.get_local_explanations().empty
+    _check_results_obj(results)
+    _check_no_skippable_files(yaml_i)
 
     yaml_i["spec"].pop("generate_explanation_files")
     yaml_i["spec"].pop("generate_forecast_file")
     yaml_i["spec"].pop("generate_metrics_file")
     operator_config = ForecastOperatorConfig.from_dict(yaml_i)
     results = operate(operator_config)
+    _check_results_obj(results)
     files = os.listdir(yaml_i["spec"]["output_directory"]["url"])
     if "errors.json" in files:
         with open(
@@ -1012,6 +1024,12 @@ def test_generate_files(operator_setup, model):
     assert "local_explanation.csv" in files, "Failed to generated local expl file"
     assert "global_explanation.csv" in files, "Failed to generated global expl file"
 
+    # Test that the results object still generates when report.html has an error
+    yaml_i["spec"]["output_directory"]["url"] = "s3://test@test/test_dir"
+    operator_config = ForecastOperatorConfig.from_dict(yaml_i)
+    results = operate(operator_config)
+    _check_results_obj(results)
+
 
 if __name__ == "__main__":
     pass