diff --git a/test_unstructured/metrics/test_evaluate.py b/test_unstructured/metrics/test_evaluate.py index 750b9ef18..478e052ad 100644 --- a/test_unstructured/metrics/test_evaluate.py +++ b/test_unstructured/metrics/test_evaluate.py @@ -1,5 +1,6 @@ import os import pathlib +import shutil import pandas as pd import pytest @@ -21,7 +22,21 @@ GOLD_CCT_DIRNAME = "gold_standard_cct" GOLD_ELEMENT_TYPE_DIRNAME = "gold_standard_element_type" +@pytest.fixture() +def _cleanup_after_test(): + # This is where the test runs + yield + + os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME) + export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct") + + # Cleanup the directory and file + if os.path.exists(export_dir): + shutil.rmtree(export_dir) + + @pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container") +@pytest.mark.usefixtures("_cleanup_after_test") def test_text_extraction_evaluation(): output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME) source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME) @@ -37,6 +52,7 @@ def test_text_extraction_evaluation(): @pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container") +@pytest.mark.usefixtures("_cleanup_after_test") def test_element_type_evaluation(): output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME) source_dir = os.path.join(TESTING_FILE_DIR, GOLD_ELEMENT_TYPE_DIRNAME) @@ -52,6 +68,7 @@ def test_element_type_evaluation(): @pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container") +@pytest.mark.usefixtures("_cleanup_after_test") def test_text_extraction_takes_list(): output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME) output_list = ["currency.csv.json"] @@ -69,6 +86,7 @@ def test_text_extraction_takes_list(): @pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container") +@pytest.mark.usefixtures("_cleanup_after_test") def test_text_extraction_grouping(): output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME) source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)