From ef67cffc0a3f932fa75ce152723362087fd53940 Mon Sep 17 00:00:00 2001
From: McKnight-42 <matthew.mcknight@dbtlabs.com>
Date: Mon, 3 Jun 2024 15:02:40 -0500
Subject: [PATCH] init push for issue 10198

---
 core/dbt/task/test.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/core/dbt/task/test.py b/core/dbt/task/test.py
index 2ae65dc3ebe..d4f8a420c27 100644
--- a/core/dbt/task/test.py
+++ b/core/dbt/task/test.py
@@ -126,6 +126,8 @@ def before_execute(self):
     def execute_data_test(self, data_test: TestNode, manifest: Manifest) -> TestResultData:
         context = generate_runtime_model_context(data_test, self.config, manifest)
 
+        hook_ctx = self.adapter.pre_model_hook(context)
+
         materialization_macro = manifest.find_materialization_macro_by_name(
             self.config.project_name, data_test.get_materialization(), self.adapter.type()
         )
@@ -140,10 +142,13 @@ def execute_data_test(self, data_test: TestNode, manifest: Manifest) -> TestResu
                 "Invalid materialization context generated, missing config: {}".format(context)
             )
 
-        # generate materialization macro
-        macro_func = MacroGenerator(materialization_macro, context)
-        # execute materialization macro
-        macro_func()
+        try:
+            # generate materialization macro
+            macro_func = MacroGenerator(materialization_macro, context)
+            # execute materialization macro
+            macro_func()
+        finally:
+            self.adapter.post_model_hook(context, hook_ctx)
         # load results from context
         # could eventually be returned directly by materialization
         result = context["load_result"]("main")
@@ -198,6 +203,8 @@ def execute_unit_test(
         # materialization, not compile the node.compiled_code
         context = generate_runtime_model_context(unit_test_node, self.config, unit_test_manifest)
 
+        hook_ctx = self.adapter.pre_model_hook(context)
+
         materialization_macro = unit_test_manifest.find_materialization_macro_by_name(
             self.config.project_name, unit_test_node.get_materialization(), self.adapter.type()
         )
@@ -213,16 +220,18 @@ def execute_unit_test(
                 "Invalid materialization context generated, missing config: {}".format(context)
             )
 
-        # generate materialization macro
-        macro_func = MacroGenerator(materialization_macro, context)
-        # execute materialization macro
         try:
+            # generate materialization macro
+            macro_func = MacroGenerator(materialization_macro, context)
+            # execute materialization macro
             macro_func()
         except DbtBaseException as e:
             raise DbtRuntimeError(
                 f"An error occurred during execution of unit test '{unit_test_def.name}'. "
                 f"There may be an error in the unit test definition: check the data types.\n {e}"
             )
+        finally:
+            self.adapter.post_model_hook(context, hook_ctx)
 
         # load results from context
         # could eventually be returned directly by materialization