diff --git a/grai-server/app/connections/tests/test_tasks.py b/grai-server/app/connections/tests/test_tasks.py
index 0ca7919fd..31d67ab8f 100644
--- a/grai-server/app/connections/tests/test_tasks.py
+++ b/grai-server/app/connections/tests/test_tasks.py
@@ -1,7 +1,7 @@
 import os
 import uuid
 from datetime import date
-from unittest import mock
+from grai_schemas.integrations.errors import NoConnectionError
 
 import pytest
 from decouple import config
@@ -217,13 +217,16 @@ def test_run_update_server_postgres_no_host(self, test_workspace, test_postgres_
         )
         run = Run.objects.create(connection=connection, workspace=test_workspace, source=test_source)
 
-        with pytest.raises(Exception) as e_info:
-            process_run(str(run.id))
+        process_run(str(run.id))
 
+        run.refresh_from_db()
+
+        assert run.status == "error"
+        assert run.metadata["error"] == "No connection"
         assert (
-            str(e_info.value)
+            run.metadata["message"]
             == 'could not translate host name "a" to address: nodename nor servname provided, or not known\n'
-            or str(e_info.value)
+            or run.metadata["message"]
             == 'could not translate host name "a" to address: Temporary failure in name resolution\n'
         )
 
@@ -784,13 +787,16 @@ def test_run_connection_schedule_postgres(self, test_workspace, test_postgres_co
             source=test_source,
         )
 
-        with pytest.raises(Exception) as e_info:
-            run_connection_schedule(str(connection.id))
+        run_connection_schedule(str(connection.id))
 
+        run = connection.runs.last()
+
+        assert run.status == "error"
+        assert run.metadata["error"] == "No connection"
         assert (
-            str(e_info.value)
+            run.metadata["message"]
             == 'could not translate host name "a" to address: nodename nor servname provided, or not known\n'
-            or str(e_info.value)
+            or run.metadata["message"]
             == 'could not translate host name "a" to address: Temporary failure in name resolution\n'
         )