fix unit test for fill_missing_col

ONSdigital · Nov 26, 2024 · d6f9915 · d6f9915
1 parent 6084901
commit d6f9915
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/tests/cdp/io/test_cdsw_output.py b/tests/cdp/io/test_cdsw_output.py
@@ -76,8 +76,21 @@ def test_insert_df_to_hive_table_with_missing_columns(
         # Mock the table's columns to include 'address'
         mock_table.return_value.columns = ["id", "name", "age", "address"]
 
+        # Mock the Hive table schema
+        mock_table_schema = T.StructType(
+            [
+                T.StructField("id", T.IntegerType()),
+                T.StructField("name", T.StringType()),
+                T.StructField("age", T.IntegerType()),
+                T.StructField(
+                    "address",
+                    T.StringType(),
+                ),  # Include 'address' with StringType
+            ],
+        )
         # Create a mock of the DataFrame (test_df) that does not contain 'address'
         test_df_mock = MagicMock()
+        mock_table.return_value.schema = mock_table_schema
         test_df_mock.columns = [
             "id",
             "name",