added tests for commented with statements

dbt-msft · Aug 23, 2024 · e8e69c7 · e8e69c7
1 parent 6fe8cf5
commit e8e69c7
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 0 deletions.
diff --git a/dbt/include/sqlserver/macros/materializations/tests.sql b/dbt/include/sqlserver/macros/materializations/tests.sql
@@ -0,0 +1,46 @@
+{% macro sqlserver__get_test_sql(main_sql, fail_calc, warn_if, error_if, limit) -%}
+
+  -- Create target schema if it does not
+  USE [{{ target.database }}];
+  IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name = '{{ target.schema }}')
+  BEGIN
+    EXEC('CREATE SCHEMA [{{ target.schema }}]')
+  END
+
+  {% set with_statement_pattern = 'with .+ as\s*\(' %}
+  {% set re = modules.re %}
+  {% set is_match = re.search(with_statement_pattern, main_sql, re.IGNORECASE) %}
+
+  {% if is_match %}
+    {% set testview %}
+      [{{ target.schema }}.testview_{{ range(1300, 19000) | random }}]
+    {% endset %}
+
+    {% set sql = main_sql.replace("'", "''")%}
+    EXEC('create view {{testview}} as {{ sql }};')
+    select
+      {{ "top (" ~ limit ~ ')' if limit != none }}
+      {{ fail_calc }} as failures,
+      case when {{ fail_calc }} {{ warn_if }}
+        then 'true' else 'false' end as should_warn,
+      case when {{ fail_calc }} {{ error_if }}
+        then 'true' else 'false' end as should_error
+    from (
+      select * from {{testview}}
+    ) dbt_internal_test;
+
+    EXEC('drop view {{testview}};')
+
+  {% else -%}
+    select
+      {{ "top (" ~ limit ~ ')' if limit != none }}
+      {{ fail_calc }} as failures,
+      case when {{ fail_calc }} {{ warn_if }}
+        then 'true' else 'false' end as should_warn,
+      case when {{ fail_calc }} {{ error_if }}
+        then 'true' else 'false' end as should_error
+    from (
+      {{ main_sql }}
+    ) dbt_internal_test
+  {%- endif -%}
+{%- endmacro %}
diff --git a/tests/functional/adapter/mssql/test_test_with.py b/tests/functional/adapter/mssql/test_test_with.py
@@ -37,6 +37,15 @@
       field: ID
 """
 
+comments_model_yml = """
+version: 2
+models:
+- name: sample_model
+  data_tests:
+  - with_statement_comments:
+      field: ID
+"""
+
 with_test_fail_sql = """
 {% test with_statement_fail(model, field) %}
 
@@ -63,6 +72,18 @@
 {% endtest %}
 """
 
+with_test_with_comments_sql = """
+{% test with_statement_comments(model, field) %}
+-- comments
+with test_sample AS (
+    SELECT {{ field }} FROM {{ model }}
+    GROUP BY {{ field }}
+    HAVING COUNT(*) > 2
+)
+SELECT * FROM test_sample
+{% endtest %}
+"""
+
 
 class BaseSQLTestWith:
     @pytest.fixture(scope="class")
@@ -77,6 +98,7 @@ def macros(self):
         return {
             "with_statement_pass.sql": with_test_pass_sql,
             "with_statement_fail.sql": with_test_fail_sql,
+            "with_statement_comments.sql": with_test_with_comments_sql,
         }
 
     @pytest.fixture(scope="class")
@@ -111,3 +133,16 @@ def models(self):
     def test_sql_test_contains_with(self, project):
         run_dbt(["run"])
         run_dbt(["test"], expect_pass=False)
+
+
+class TestSQLTestWithComment(BaseSQLTestWith):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "sample_model.sql": sample_model,
+            "schema.yml": comments_model_yml,
+        }
+
+    def test_sql_test_contains_with(self, project):
+        run_dbt(["run"])
+        run_dbt(["test"])