diff --git a/docs/gallery/data_testing/data_testing_flow.py b/docs/gallery/data_testing/data_testing_flow.py
index 38b771e..dd39164 100644
--- a/docs/gallery/data_testing/data_testing_flow.py
+++ b/docs/gallery/data_testing/data_testing_flow.py
@@ -110,7 +110,7 @@ with Flow(
 
     # ↓↓↓ This is where we actually run the data test. The result of the test
     # gets stored in data_test_result.
-    data_test_result = run_data_test_on_mysql.run(
+    data_test_result = run_data_test_on_mysql(
         name=DATA_TEST_NAME,  # ←←← The name we set earlier
         # ↓↓↓ The credentials to the MySQL where the data lives. We pass the
         # ssh tunnel host and port instead of the true MySQL because we want
@@ -126,7 +126,7 @@ with Flow(
         },  # in the queries regardless.
         query=DATA_TEST_QUERY,  # ←←← The query we set earlier
         expectation_configurations=DATA_TEST_EXPECTATIONS,  # ←←← Idem
-        upstream_tasks=[ssh_tunnel] # ←←← We must wait for the tunnel to be ready
+        upstream_tasks=[ssh_tunnel],  # ←←← We must wait for the tunnel to be ready
     )
 
     # ↑↑↑ will take care of everything: connecting to S3 and DW, generate all
@@ -141,8 +141,7 @@ with Flow(
     # The following snippets are optional. You should judge if you want to do
     # something similar or not in your flow based on your needs.
 
-    tunnel_closed = close_ssh_tunnel(ssh_tunnel)
-
+    tunnel_closed = close_ssh_tunnel(ssh_tunnel, upstream_tasks=[data_test_result])
 
 
 # TODO
diff --git a/docs/gallery/data_testing/requirements.txt b/docs/gallery/data_testing/requirements.txt
new file mode 100644
index 0000000..245292b
--- /dev/null
+++ b/docs/gallery/data_testing/requirements.txt
@@ -0,0 +1,4 @@
+prefect==1.2.2
+great_expectations==0.15.45
+SQLAlchemy==1.4.46
+lolafect==0.4.0
\ No newline at end of file