Refactor row count tests

calogica · Nov 2, 2020 · 4b92aa9 · 4b92aa9
1 parent c6125e9
commit 4b92aa9
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 17 deletions.
diff --git a/macros/schema_tests/table_shape/expect_table_row_count_to_be_between.sql b/macros/schema_tests/table_shape/expect_table_row_count_to_be_between.sql
@@ -4,16 +4,14 @@
                                                       partition_column=None,
                                                       partition_filter=None
                                                     ) -%}
-with row_count as (
-    select count(*) as cnt 
-    from {{ model }}
-    {% if partition_column and partition_filter %}
-    where {{ partition_column }} {{ partition_filter }}
-    {% endif %}
-)
-select count(*)
-from 
-    row_count
-where cnt < {{ minimum }} or
-      cnt > {{ maximum }}
+{% set expression %}
+count(*) 
+{% endset %}
+{{ dbt_expectations.expression_between(model, 
+                                        expression=expression,
+                                        minimum=minimum, 
+                                        maximum=maximum, 
+                                        partition_column=partition_column, 
+                                        partition_filter=partition_filter
+                                        ) }}      
 {%- endmacro -%}
diff --git a/macros/schema_tests/table_shape/expect_table_row_count_to_equal.sql b/macros/schema_tests/table_shape/expect_table_row_count_to_equal.sql
@@ -3,9 +3,12 @@
                                                 partition_column=None,
                                                 partition_filter=None
                                                 ) -%}
-select abs({{ expected_number_of_rows }} - count(*)) 
-from {{ model }}
-{% if partition_column and partition_filter %}
-where {{ partition_column }} {{ partition_filter }}
-{% endif %}
+{% set filter_cond = partition_column ~ " " ~ partition_filter if partition_column and partition_filter else None %}
+{% set expression %}
+count(*) = {{ expected_number_of_rows }}
+{% endset %}
+{{ dbt_expectations.expression_is_true(model, 
+                                        expression=expression,
+                                        filter_cond=filter_cond)
+                                        }}
 {%- endmacro -%}