Fix sum_total in expect_multicolumn_sum_to_equal to accept column nam…

…es (#291) * add sum around sum total so columns can be used * add col_numeric_a_plus_b to data_test add test for column name in expect_multicolumn_sum_to_equal * add a test to get row level matches * Update dbt-spark version param --------- Co-authored-by: Claus Herther <[email protected]>
calogica · Nov 2, 2023 · bf3f7fe · bf3f7fe
1 parent 009a831
commit bf3f7fe
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 10 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -119,7 +119,7 @@ jobs:
           command: |
             python3 -m venv venv
             . venv/bin/activate
-            pip install dbt-spark "dbt-spark[PyHive]"
+            pip install "dbt-spark==$DBT_VERSION" "dbt-spark[PyHive]==$DBT_VERSION"
       - run: *dbt-deps
       - run:
           name: Wait for Spark-Thrift

diff --git a/integration_tests/models/schema_tests/data_test.sql b/integration_tests/models/schema_tests/data_test.sql
@@ -6,7 +6,8 @@ select
     'a' as col_string_a,
     'b' as col_string_b,
     cast(null as {{ dbt.type_string() }}) as col_null,
-    cast(null as {{ dbt.type_string() }}) as col_null_2
+    cast(null as {{ dbt.type_string() }}) as col_null_2,
+    1.0 as col_numeric_a_plus_b
 
 union all
 
@@ -18,7 +19,8 @@ select
     'b' as col_string_a,
     'ab' as col_string_b,
     null as col_null,
-    null as col_null_2
+    null as col_null_2,
+    1.0 as col_numeric_a_plus_b
 
 union all
 
@@ -30,7 +32,8 @@ select
     'c' as col_string_a,
     'abc' as col_string_b,
     null as col_null,
-    null as col_null_2
+    null as col_null_2,
+    1.0 as col_numeric_a_plus_b
 
 union all
 
@@ -42,4 +45,5 @@ select
     'c' as col_string_a,
     'abcd' as col_string_b,
     null as col_null,
-    null as col_null_2
+    null as col_null_2,
+    1.0 as col_numeric_a_plus_b
diff --git a/integration_tests/models/schema_tests/schema.yml b/integration_tests/models/schema_tests/schema.yml
@@ -417,7 +417,7 @@ models:
           row_condition: 1=1
           compare_row_condition: 1=1
       - dbt_expectations.expect_table_column_count_to_equal:
-          value: 8
+          value: 9
       - dbt_expectations.expect_table_column_count_to_be_between:
           min_value: 1
           max_value: 10
@@ -428,9 +428,9 @@ models:
       - dbt_expectations.expect_table_columns_to_contain_set:
           column_list: ["col_numeric_b", "col_string_a"]
       - dbt_expectations.expect_table_columns_to_match_set:
-          column_list: ["idx", "date_col", "col_numeric_a", "col_numeric_b", "col_string_a", "col_string_b", "col_null", "col_null_2"]
+          column_list: ["idx", "date_col", "col_numeric_a", "col_numeric_b", "col_string_a", "col_string_b", "col_null", "col_null_2", "col_numeric_a_plus_b"]
       - dbt_expectations.expect_table_columns_to_match_ordered_list:
-          column_list: ["idx", "date_col", "col_numeric_a", "col_numeric_b", "col_string_a", "col_string_b", "col_null", "col_null_2"]
+          column_list: ["idx", "date_col", "col_numeric_a", "col_numeric_b", "col_string_a", "col_string_b", "col_null", "col_null_2", "col_numeric_a_plus_b"]
       - dbt_expectations.expect_table_column_count_to_equal_other_table:
           compare_model: ref("data_test")
       - dbt_expectations.expect_table_columns_to_not_contain_set:
@@ -477,7 +477,13 @@ models:
       - dbt_expectations.expect_multicolumn_sum_to_equal:
           column_list: ["col_numeric_a", "col_numeric_b"]
           sum_total: 4
-
+      - dbt_expectations.expect_multicolumn_sum_to_equal:
+          column_list: ["col_numeric_a", "col_numeric_b"]
+          sum_total: col_numeric_a_plus_b
+      - dbt_expectations.expect_multicolumn_sum_to_equal:
+          column_list: ["col_numeric_a", "col_numeric_b"]
+          sum_total: col_numeric_a_plus_b
+          group_by: [idx]
 
 
     columns:

diff --git a/macros/schema_tests/multi-column/expect_multicolumn_sum_to_equal.sql b/macros/schema_tests/multi-column/expect_multicolumn_sum_to_equal.sql
@@ -9,7 +9,8 @@
 {% set expression %}
 {% for column in column_list %}
 sum({{ column }}){% if not loop.last %} + {% endif %}
-{% endfor %} = {{ sum_total }}
+{# the if just allows for column names or literal numbers #}
+{% endfor %} = {% if sum_total is number %}{{sum_total}}{% else %}sum({{ sum_total }}){% endif %}
 {% endset %}
 
 {{ dbt_expectations.expression_is_true(model,