Adding unit test for string flag column

ONSdigital · May 17, 2024 · e24f451 · e24f451
1 parent 0607562
commit e24f451
Showing 1 changed file with 9 additions and 2 deletions.
diff --git 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)