From db38042e86180ddaf8ede79529f2d9475da88d65 Mon Sep 17 00:00:00 2001
From: Lukas Gust <lgust@letsonboard.com>
Date: Tue, 5 Sep 2023 11:02:21 -0600
Subject: [PATCH] Added additional test

---
 tests/functional/adapter/test_partitions.py | 95 ++++++++++++++++++++-
 1 file changed, 92 insertions(+), 3 deletions(-)

diff --git a/tests/functional/adapter/test_partitions.py b/tests/functional/adapter/test_partitions.py
index 9bd91de9..f5f1e6d3 100644
--- a/tests/functional/adapter/test_partitions.py
+++ b/tests/functional/adapter/test_partitions.py
@@ -17,7 +17,7 @@
 cross join unnest(date_array) as t2(date_column)
 """
 
-test_null_valued_partitions_model_sql = """
+test_single_nullable_partition_model_sql = """
 with data as (
     select
         random() as col_1,
@@ -35,6 +35,49 @@
 select random() as col_1, NULL as id
 """
 
+test_nullable_partitions_model_sql = """
+{{ config(
+    materialized='table',
+    format='parquet',
+    s3_data_naming='table',
+    partitioned_by=['id', 'date_column']
+) }}
+
+with data as (
+    select
+    random() as rnd,
+    row_number() over() as id,
+    cast(date_column as date) as date_column
+from (
+    values (
+        sequence(from_iso8601_date('2023-01-01'), from_iso8601_date('2023-07-31'), interval '1' day)
+    )
+) as t1(date_array)
+cross join unnest(date_array) as t2(date_column)
+)
+
+select
+    rnd,
+    case when id  <= 50 then null else id end as id,
+    date_column
+from data
+union all
+select
+    random() as rnd,
+    NULL as id,
+    NULL as date_column
+union all
+select
+    random() as rnd,
+    NULL as id,
+    cast('2023-09-02' as date) as date_column
+union all
+select
+    random() as rnd,
+    40 as id,
+    NULL as date_column
+"""
+
 
 class TestHiveTablePartitions:
     @pytest.fixture(scope="class")
@@ -146,6 +189,52 @@ def test__check_incremental_run_with_partitions(self, project):
 
 
 class TestHiveNullValuedPartitions:
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "models": {
+                "+table_type": "hive",
+                "+materialized": "table",
+                "+partitioned_by": ["id", "date_column"],
+            }
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "test_nullable_partitions_model.sql": test_nullable_partitions_model_sql,
+        }
+
+    def test__check_run_with_partitions(self, project):
+        relation_name = "test_nullable_partitions_model"
+        model_run_result_row_count_query = f"select count(*) as records from {project.test_schema}.{relation_name}"
+        model_run_result_null_id_count_query = (
+            f"select count(*) as records from {project.test_schema}.{relation_name} where id is null"
+        )
+        model_run_result_null_date_count_query = (
+            f"select count(*) as records from {project.test_schema}.{relation_name} where date_column is null"
+        )
+
+        first_model_run = run_dbt(["run", "--select", relation_name])
+        first_model_run_result = first_model_run.results[0]
+
+        # check that the model run successfully
+        assert first_model_run_result.status == RunStatus.Success
+
+        records_count_first_run = project.run_sql(model_run_result_row_count_query, fetch="all")[0][0]
+
+        assert records_count_first_run == 215
+
+        null_id_count_first_run = project.run_sql(model_run_result_null_id_count_query, fetch="all")[0][0]
+
+        assert null_id_count_first_run == 52
+
+        null_date_count_first_run = project.run_sql(model_run_result_null_date_count_query, fetch="all")[0][0]
+
+        assert null_date_count_first_run == 2
+
+
+class TestHiveSingleNullValuedPartition:
     @pytest.fixture(scope="class")
     def project_config_update(self):
         return {
@@ -159,11 +248,11 @@ def project_config_update(self):
     @pytest.fixture(scope="class")
     def models(self):
         return {
-            "test_hive_partitions_null_values.sql": test_null_valued_partitions_model_sql,
+            "test_single_nullable_partition_model.sql": test_single_nullable_partition_model_sql,
         }
 
     def test__check_run_with_partitions(self, project):
-        relation_name = "test_hive_partitions_null_values"
+        relation_name = "test_single_nullable_partition_model"
         model_run_result_row_count_query = f"select count(*) as records from {project.test_schema}.{relation_name}"
 
         first_model_run = run_dbt(["run", "--select", relation_name])