diff --git a/test/test_initial.py b/test/test_initial.py index 9987aa8..1b63d36 100644 --- a/test/test_initial.py +++ b/test/test_initial.py @@ -32,6 +32,15 @@ def acro() -> ACRO: return ACRO(suppress=True) +def clean_up(name): + """Removes unwanted files or directory.""" + if os.path.exists(name): + if os.path.isfile(name): + os.remove(name) + elif os.path.isdir(name): + shutil.rmtree(name) + + def test_crosstab_without_suppression(data): """Crosstab threshold without automatic suppression.""" acro = ACRO(suppress=False) @@ -849,3 +858,10 @@ def test_crosstab_with_manual_totals_with_suppression_with_two_aggfunc( "We can not calculate the margins with a list of aggregation functions. " "Please create a table for each aggregation function" in caplog.text ) + + +def test_cleanup(): + """Gets rid of files created during tests.""" + names = ["test_outputs", "test_add_to_acro", "sdc_results", "RES_PYTEST"] + for name in names: + clean_up(name) diff --git a/test/test_stata_interface.py b/test/test_stata_interface.py index 5cc3253..561b15e 100644 --- a/test/test_stata_interface.py +++ b/test/test_stata_interface.py @@ -1,6 +1,7 @@ """This module contains unit tests for the stata interface.""" import os +import shutil import numpy as np import pandas as pd @@ -33,6 +34,15 @@ def data() -> pd.DataFrame: return data +def clean_up(name): + """Removes unwanted files or directory.""" + if os.path.exists(name): + if os.path.isfile(name): + os.remove(name) + elif os.path.isdir(name): + shutil.rmtree(name) + + def dummy_acrohandler( data, command, varlist, exclusion, exp, weights, options ): # pylint:disable=too-many-arguments @@ -736,3 +746,10 @@ def test_stata_unknown(data): ) correct = "acro command not recognised: foo" assert ret == correct, f"got:\n{ret}\nexpected:\n{correct}\n" + + +def test_cleanup(): + """Gets rid of files created during tests.""" + names = ["test_outputs"] + for name in names: + clean_up(name)