From e24f4513838bd16984dab0de51d122657c2f475b Mon Sep 17 00:00:00 2001 From: dayj1 Date: Fri, 17 May 2024 12:02:40 +0100 Subject: [PATCH] Adding unit test for string flag column --- tests/test_imputation_flags.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/test_imputation_flags.py b/tests/test_imputation_flags.py index d78bb525..a1455400 100644 --- a/tests/test_imputation_flags.py +++ b/tests/test_imputation_flags.py @@ -4,7 +4,7 @@ from helper_functions import load_and_format from pandas.testing import assert_frame_equal -from src.imputation_flags import create_impute_flags +from src.imputation_flags import create_impute_flags, generate_imputation_flag_string @pytest.fixture(scope="class") @@ -14,7 +14,7 @@ def imputation_flag_test_data(): class TestImputationFlags: def test_create_impute_flags(self, imputation_flag_test_data): - df_expected_output = imputation_flag_test_data + df_expected_output = imputation_flag_test_data.copy() df_expected_output.drop(["imputation_flag"], axis=1, inplace=True) df_input = df_expected_output.copy() df_input = df_input[ @@ -36,3 +36,10 @@ def test_create_impute_flags(self, imputation_flag_test_data): auxiliary="auxiliary", ) assert_frame_equal(df_output, df_expected_output) + + def test_imputation_flag_strings(self, imputation_flag_test_data): + df_expected_output = imputation_flag_test_data.copy() + df_input = imputation_flag_test_data.copy() + df_input.drop("imputation_flag", axis=1, inplace=True) + df_output = generate_imputation_flag_string(df_input) + assert_frame_equal(df_output, df_expected_output)