From 52e2e8531eeb9a4218cb4ac5c73fac89d0eebf8b Mon Sep 17 00:00:00 2001
From: Michelle Ark <michelle.ark@dbtlabs.com>
Date: Fri, 7 Jun 2024 12:14:34 -0400
Subject: [PATCH] test adapter pre hooks for unit testing

---
 tests/functional/unit_testing/fixtures.py     | 17 +++++++
 .../unit_testing/test_ut_adapter_hooks.py     | 49 +++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 tests/functional/unit_testing/test_ut_adapter_hooks.py

diff --git a/tests/functional/unit_testing/fixtures.py b/tests/functional/unit_testing/fixtures.py
index 3028e0bc1e6..e73351f89d8 100644
--- a/tests/functional/unit_testing/fixtures.py
+++ b/tests/functional/unit_testing/fixtures.py
@@ -116,6 +116,23 @@
         tags: test_this
 """
 
+test_my_model_pass_yml = """
+unit_tests:
+  - name: test_my_model
+    model: my_model
+    given:
+      - input: ref('my_model_a')
+        rows:
+          - {id: 1, a: 1}
+      - input: ref('my_model_b')
+        rows:
+          - {id: 1, b: 2}
+          - {id: 2, b: 2}
+    expect:
+      rows:
+        - {c: 3}
+"""
+
 
 test_my_model_simple_fixture_yml = """
 unit_tests:
diff --git a/tests/functional/unit_testing/test_ut_adapter_hooks.py b/tests/functional/unit_testing/test_ut_adapter_hooks.py
new file mode 100644
index 00000000000..48828f8b45a
--- /dev/null
+++ b/tests/functional/unit_testing/test_ut_adapter_hooks.py
@@ -0,0 +1,49 @@
+from unittest import mock
+
+import pytest
+
+from dbt.tests.util import run_dbt
+from dbt_common.exceptions import CompilationError
+from tests.functional.unit_testing.fixtures import (
+    my_model_a_sql,
+    my_model_b_sql,
+    my_model_sql,
+    test_my_model_pass_yml,
+)
+
+
+class BaseUnitTestAdapterHook:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "my_model.sql": my_model_sql,
+            "my_model_a.sql": my_model_a_sql,
+            "my_model_b.sql": my_model_b_sql,
+            "test_my_model.yml": test_my_model_pass_yml,
+        }
+
+
+class TestUnitTestAdapterHookPasses(BaseUnitTestAdapterHook):
+    def test_unit_test_runs_adapter_pre_hook(self, project):
+        results = run_dbt(["run"])
+        assert len(results) == 3
+
+        mock_pre_model_hook = mock.Mock()
+        with mock.patch.object(type(project.adapter), "pre_model_hook", mock_pre_model_hook):
+            results = run_dbt(["test", "--select", "test_name:test_my_model"], expect_pass=True)
+
+            assert len(results) == 1
+            mock_pre_model_hook.assert_called_once()
+
+
+class TestUnitTestAdapterHookFails(BaseUnitTestAdapterHook):
+    def test_unit_test_runs_adapter_pre_hook_fails(self, project):
+        results = run_dbt(["run"])
+        assert len(results) == 3
+
+        mock_pre_model_hook = mock.Mock()
+        mock_pre_model_hook.side_effect = CompilationError("exception from adapter.pre_model_hook")
+        with mock.patch.object(type(project.adapter), "pre_model_hook", mock_pre_model_hook):
+            run_dbt(["test", "--select", "test_name:test_my_model"], expect_pass=False)
+
+            mock_pre_model_hook.assert_called_once()