Add support for struct/array columns and allow udf definitions to be set

integration-tests/dbt_project.yml

@@ -33,3 +33,5 @@ vars:
     type_separator: "::"
     diff_column: "diff"
     count_column: "count"
+    convert_columns: []
+    udf_definitions: ""

macros/tests.sql

@@ -78,20 +78,26 @@
   {% set expectations = test_configuration.expectations %}
   {% set model_node = dbt_unit_testing.model_node(test_configuration.model_node) %}
   {%- set model_complete_sql = dbt_unit_testing.build_model_complete_sql(model_node, test_configuration.mocks, test_configuration.options) -%}
-  {% set columns = dbt_unit_testing.quote_and_join_columns(dbt_unit_testing.extract_columns_list(expectations.input_values)) %}
+  {% set convert_columns = test_configuration.options.convert_columns | default("convert_columns") %}
+  {% set select_columns = dbt_unit_testing.quote_and_join_columns(dbt_unit_testing.extract_columns_list(expectations.input_values), convert_columns, add_alias=true) %}
+  {% set group_columns = dbt_unit_testing.quote_and_join_columns(dbt_unit_testing.extract_columns_list(expectations.input_values), convert_columns) %}
+  {% set result_columns = dbt_unit_testing.quote_and_join_columns(dbt_unit_testing.extract_columns_list(expectations.input_values)) %}
 
   {% set diff_column = test_configuration.options.diff_column | default("diff") %}
   {% set count_column = test_configuration.options.count_column | default("count") %}
+  {% set udf_definitions = test_configuration.options.udf_definitions | default("udf_definitions") %}
 
   {%- set actual_query -%}
-    select count(1) as {{ count_column }}, {{columns}} from ( {{ model_complete_sql }} ) as s group by {{ columns }}
+    select count(1) as {{ count_column }}, {{ select_columns }} from ( {{ model_complete_sql }} ) as s group by {{ group_columns }}
   {% endset %}
 
   {%- set expectations_query -%}
-    select count(1) as {{ count_column }}, {{columns}} from ({{ expectations.input_values }}) as s group by {{ columns }}
+    select count(1) as {{ count_column }}, {{ select_columns }} from ({{ expectations.input_values }}) as s group by {{ group_columns }}
   {% endset %}
 
   {%- set test_query -%}
+    {{ udf_definitions }}
+
     with expectations as (
       {{ expectations_query }}
     ),
@@ -100,15 +106,15 @@
     ),
 
     extra_entries as (
-    select '+' as {{ diff_column }}, {{ count_column }}, {{columns}} from actual
+    select '+' as {{ diff_column }}, {{ count_column }}, {{ result_columns }} from actual
     {{ except() }}
-    select '+' as {{ diff_column }}, {{ count_column }}, {{columns}} from expectations),
+    select '+' as {{ diff_column }}, {{ count_column }}, {{ result_columns }} from expectations),
 
     missing_entries as (
-    select '-' as {{ diff_column }}, {{ count_column }}, {{columns}} from expectations
+    select '-' as {{ diff_column }}, {{ count_column }}, {{ result_columns }} from expectations
     {{ except() }}
-    select '-' as {{ diff_column }}, {{ count_column }}, {{columns}} from actual)
-    
+    select '-' as {{ diff_column }}, {{ count_column }}, {{ result_columns }} from actual)
+
     select * from extra_entries
     UNION ALL
     select * from missing_entries
@@ -150,9 +156,12 @@
   {% set actual_query = test_queries.actual_query %}
   {% set expectations_query = test_queries.expectations_query %}
   {% set test_query = test_queries.test_query %}
+  {% set udf_definitions = test_configuration.options.udf_definitions | default("udf_definitions") %}
 
   {%- set count_query -%}
-    select * FROM 
+    {{ udf_definitions }}
+
+    select * FROM
       (select count(1) as expectation_count from (
         {{ expectations_query }}
       ) as exp) as exp_count,

macros/utils.sql

@@ -23,9 +23,20 @@
   {{ return(columns) }}
 {% endmacro %}
 
-{% macro quote_and_join_columns(columns) %}
-  {% set columns = dbt_unit_testing.map(columns, dbt_unit_testing.quote_identifier) | join(",") %}
-  {{ return (columns) }}
+{% macro quote_and_join_columns(columns, convert_columns, add_alias = false) %}
+    {% set result = [] %}
+    {% for column in columns %}
+        {% if column in convert_columns %}
+            {% set modified_column = "TO_JSON_STRING(" ~ dbt_unit_testing.quote_identifier(column) ~ ")" %}
+            {% if add_alias %}
+                {% set modified_column = modified_column ~ " AS " ~ dbt_unit_testing.quote_identifier(column) %}
+            {% endif %}
+        {% else %}
+            {% set modified_column = dbt_unit_testing.quote_identifier(column) %}
+        {% endif %}
+        {% do result.append(modified_column) %}
+    {% endfor %}
+    {{ return(result | join(", ")) }}
 {% endmacro %}
 
 {% macro sql_encode(s) %}