tests for flatten json

tests/warehouse/test_bigquery_ops.py

@@ -7,6 +7,7 @@
 import subprocess
 from logging import basicConfig, getLogger, INFO
 from dbt_automation.operations.droprenamecolumns import rename_columns, drop_columns
+from dbt_automation.operations.flattenjson import flattenjson
 from dbt_automation.operations.mergeoperations import (
     merge_operations,
 )
@@ -945,3 +946,45 @@ def test_merge_operation(self):
             )
             == 0
         )
+
+    def test_flattenjson(self):
+        """Test flattenjson."""
+        wc_client = TestBigqueryOperations.wc_client
+        output_name = "flatten_json"
+
+        config = {
+            "input": {
+                "input_type": "source",
+                "input_name": "_airbyte_raw_Sheet1",
+                "source_name": "sample",
+            },
+            "source_schema": "pytest_intermediate",
+            "dest_schema": "pytest_intermediate",
+            "output_name": output_name,
+            "source_columns": [
+                "_airbyte_ab_id",
+                "_airbyte_data",
+                "_airbyte_emitted_at",
+            ],
+            "json_column": "_airbyte_data",
+            "json_columns_to_copy": ["NGO", "SPOC", "Month", "measure1", "measure2", "Indicator"],
+        }
+
+        flattenjson(config, wc_client, TestBigqueryOperations.test_project_dir)
+
+        TestBigqueryOperations.execute_dbt("run", output_name)
+
+        cols = [
+            col_dict["name"]
+            for col_dict in wc_client.get_table_columns(
+                "pytest_intermediate", output_name
+            )
+        ]
+
+        assert "_airbyte_data_NGO" in cols
+        assert "_airbyte_data_Month" in cols
+        assert "_airbyte_ab_id" in cols
+
+        table_data = wc_client.get_table_data("pytest_intermediate", output_name, 1)
+        assert table_data[0]["_airbyte_data_NGO"] == 'BAMANEH'
+        assert table_data[0]["_airbyte_data_SPOC"] == 'SPOC A'

tests/warehouse/test_postgres_ops.py

@@ -742,13 +742,13 @@ def test_mergetables(self):
     def test_flattenjson(self):
         """Test flattenjson."""
         wc_client = TestPostgresOperations.wc_client
-        output_name = "flatten"
+        output_name = "flatten_json"
 
         config = {
             "input": {
-                "input_type": "model",
+                "input_type": "source",
                 "input_name": "_airbyte_raw_Sheet1",
-                "source_name": None,
+                "source_name": "sample",
             },
             "dest_schema": "pytest_intermediate",
             "source_schema": "pytest_intermediate",
@@ -772,10 +772,13 @@ def test_flattenjson(self):
                 "pytest_intermediate", output_name
             )
         ]
-        assert "output_column_name" in cols
+        assert "_airbyte_data_NGO" in cols
+        assert "_airbyte_data_Month" in cols
+        assert "_airbyte_ab_id" in cols
 
         table_data = wc_client.get_table_data("pytest_intermediate", output_name, 1)
-        assert table_data[0]["output_column_name"] == "NGO"  # Replace with expected value
+        assert table_data[0]["_airbyte_data_NGO"] == 'BAMANEH'
+        assert table_data[0]["_airbyte_data_SPOC"] == 'SPOC A'
 
     def test_merge_operation(self):
         """test merge_operation"""