updated behaviour of custom tests using with statements

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 %}

tests/functional/adapter/mssql/test_test_with.py

@@ -0,0 +1,148 @@
+import pytest
+from dbt.tests.util import run_dbt
+
+sample_model = """
+SELECT
+    1 as ID,
+    'a' as data
+
+UNION ALL
+
+SELECT
+    2 as ID,
+    'b' as data
+
+UNION ALL
+
+SELECT
+    2 as ID,
+    'c' as data
+"""
+
+pass_model_yml = """
+version: 2
+models:
+- name: sample_model
+  data_tests:
+  - with_statement_pass:
+      field: ID
+"""
+
+fail_model_yml = """
+version: 2
+models:
+- name: sample_model
+  data_tests:
+  - with_statement_fail:
+      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) %}
+
+with test_sample AS (
+    SELECT {{ field }} FROM {{ model }}
+    GROUP BY {{ field }}
+    HAVING COUNT(*) > 1
+)
+SELECT * FROM test_sample
+
+{% endtest %}
+"""
+
+with_test_pass_sql = """
+{% test with_statement_pass(model, field) %}
+
+with test_sample AS (
+    SELECT {{ field }} FROM {{ model }}
+    GROUP BY {{ field }}
+    HAVING COUNT(*) > 2
+)
+SELECT * FROM test_sample
+
+{% 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")
+    def project_config_update(self):
+        return {
+            "config-version": 2,
+            "macro-paths": ["macros"],
+        }
+
+    @pytest.fixture(scope="class")
+    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")
+    def models(self):
+        return {
+            "sample_model.sql": sample_model,
+            "schema.yml": pass_model_yml,
+        }
+
+
+class TestSQLTestWithPass(BaseSQLTestWith):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "sample_model.sql": sample_model,
+            "schema.yml": pass_model_yml,
+        }
+
+    def test_sql_test_contains_with(self, project):
+        run_dbt(["run"])
+        run_dbt(["test"])
+
+
+class TestSQLTestWithFail(BaseSQLTestWith):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "sample_model.sql": sample_model,
+            "schema.yml": fail_model_yml,
+        }
+
+    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"])