add table level test

Signed-off-by: jaga-sknt <[email protected]>
ubie-oss · Oct 15, 2024 · 0b15945 · 0b15945
1 parent 71e5050
commit 0b15945
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 3 deletions.
diff --git a/integration_tests/macros/tests/codegen/test_generate_secured_model_schema_v2.sql b/integration_tests/macros/tests/codegen/test_generate_secured_model_schema_v2.sql
@@ -72,6 +72,21 @@
       labels={
         "key1": "value1",
         "key2": "value2",
+      },
+      data_tests={
+        "test1": {
+          "key1": "value1",
+        },
+        "test2": {
+          "key2": "value2",
+          "key3": "value3",
+        }
+      },
+      tests={
+        "test1": {
+          "key4": "value4",
+          "key5": "value5",
+        }
       }
     ) -%}
 
@@ -86,6 +101,15 @@ models:
     description: |-
       Sample description
     tags: ['tag1']
+    data_tests:
+      - test1:
+          key1: value1
+      - test2:
+          key2: value2
+          key3: value3
+      - test1:
+          key4: value4
+          key5: value5
     meta:
       key1: value1
       key2: value2

diff --git a/macros/codegen/generate_secured_model_schema_v2.sql b/macros/codegen/generate_secured_model_schema_v2.sql
@@ -7,7 +7,9 @@
       description=none,
       columns={},
       tags=[],
-      labels={}
+      labels={},
+      data_tests={},
+      tests={}
     ) %}
   {{- return(adapter.dispatch('generate_secured_model_schema_v2', 'dbt_data_privacy')(
       objective=objective,
@@ -18,7 +20,9 @@
       description=description,
       columns=columns,
       tags=tags,
-      labels=labels
+      labels=labels,
+      data_tests=data_tests,
+      tests=tests
     )) -}}
 {% endmacro %}
 
@@ -31,7 +35,9 @@
       description=none,
       columns=[],
       tags=[],
-      labels={}
+      labels={},
+      data_tests={},
+      tests={}
     ) %}
 
   {% if columns | length == 0 %}
@@ -58,6 +64,26 @@ models:
     tags: {{ tags | unique | sort | list }}
     {%- endif %}
     {%- if labels | length > 0 %}
+    {#- We support both `data_tests` and `tests` for a while -#}
+    {#- TODO: remove `tests` after dbt no longer supports it -#}
+    {%- set combined_tests = [] %}
+    {%- for test_name, test_content in data_tests.items() %}
+        {%- do combined_tests.append({test_name: [test_content]}) %}
+    {%- endfor %}
+    {%- for test_name, test_content in tests.items() %}
+        {%- do combined_tests.append({test_name: [test_content]}) %}
+    {%- endfor %}
+    {%- if combined_tests | length > 0 %}
+    data_tests:
+        {%- for test in combined_tests %}
+        - {{ test.keys() | first }}:
+            {%- for content in test.values() | first %}
+            {%- for key, value in content.items() %}
+            {{ key }}: {{ value }}
+            {%- endfor %}
+            {%- endfor %}
+        {%- endfor %}
+    {%- endif %}
     meta: {%- for k, v in labels.items() %}
       {{ k }}: {{ v }}
     {%- endfor %}