Fixes tests

PrefectHQ · Jan 31, 2025 · a4c97c4 · a4c97c4
1 parent f4c8bb8
commit a4c97c4
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 28 deletions.
diff --git a/tests/client/test_prefect_client.py b/tests/client/test_prefect_client.py
@@ -2872,7 +2872,7 @@ async def test_worker_heartbeat_sends_metadata_if_passed(
         self, prefect_client: PrefectClient
     ):
         with mock.patch(
-            "prefect.client.orchestration.PrefectHttpxAsyncClient.post",
+            "prefect.client.orchestration.base.BaseAsyncClient.request",
             return_value=httpx.Response(status_code=204),
         ) as mock_post:
             await prefect_client.send_worker_heartbeat(
@@ -2895,7 +2895,7 @@ async def test_worker_heartbeat_does_not_send_metadata_if_not_passed(
         self, prefect_client: PrefectClient
     ):
         with mock.patch(
-            "prefect.client.orchestration.PrefectHttpxAsyncClient.post",
+            "prefect.client.orchestration.base.BaseAsyncClient.request",
             return_value=httpx.Response(status_code=204),
         ) as mock_post:
             await prefect_client.send_worker_heartbeat(

diff --git a/tests/workers/test_base_worker.py b/tests/workers/test_base_worker.py
@@ -16,7 +16,7 @@
 from prefect.client.base import ServerType
 from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas import FlowRun
-from prefect.client.schemas.objects import StateType, WorkerMetadata
+from prefect.client.schemas.objects import Integration, StateType, WorkerMetadata
 from prefect.exceptions import (
     CrashedRun,
     ObjectNotFound,
@@ -1990,7 +1990,7 @@ async def test_worker_heartbeat_sends_integrations(
                     "prefect.workers.base.load_prefect_collections"
                 ) as mock_load_prefect_collections,
                 mock.patch(
-                    "prefect.client.orchestration.PrefectHttpxAsyncClient.post"
+                    "prefect.client.orchestration._work_pools.client.WorkPoolAsyncClient.send_worker_heartbeat",
                 ) as mock_send_worker_heartbeat_post,
                 mock.patch("prefect.workers.base.distributions") as mock_distributions,
             ):
@@ -2010,17 +2010,13 @@ async def test_worker_heartbeat_sends_integrations(
                     await worker.sync_with_backend()
 
                 mock_send_worker_heartbeat_post.assert_called_once_with(
-                    f"/work_pools/{work_pool.name}/workers/heartbeat",
-                    json={
-                        "name": worker.name,
-                        "heartbeat_interval_seconds": worker.heartbeat_interval_seconds,
-                        "metadata": {
-                            "integrations": [
-                                {"name": "prefect-aws", "version": "1.0.0"}
-                            ]
-                        },
-                        "return_id": True,
-                    },
+                    work_pool_name=work_pool.name,
+                    worker_name=worker.name,
+                    heartbeat_interval_seconds=30.0,
+                    get_worker_id=True,
+                    worker_metadata=WorkerMetadata(
+                        integrations=[Integration(name="prefect-aws", version="1.0.0")]
+                    ),
                 )
 
             assert worker._worker_metadata_sent
@@ -2050,7 +2046,7 @@ async def _worker_metadata(self) -> WorkerMetadata:
                     "prefect.workers.base.load_prefect_collections"
                 ) as mock_load_prefect_collections,
                 mock.patch(
-                    "prefect.client.orchestration.PrefectHttpxAsyncClient.post"
+                    "prefect.client.orchestration._work_pools.client.WorkPoolAsyncClient.send_worker_heartbeat",
                 ) as mock_send_worker_heartbeat_post,
                 mock.patch("prefect.workers.base.distributions") as mock_distributions,
             ):
@@ -2070,18 +2066,14 @@ async def _worker_metadata(self) -> WorkerMetadata:
                     await worker.sync_with_backend()
 
                 mock_send_worker_heartbeat_post.assert_called_once_with(
-                    f"/work_pools/{work_pool.name}/workers/heartbeat",
-                    json={
-                        "name": worker.name,
-                        "heartbeat_interval_seconds": worker.heartbeat_interval_seconds,
-                        "metadata": {
-                            "integrations": [
-                                {"name": "prefect-aws", "version": "1.0.0"}
-                            ],
-                            "custom_field": "heya",
-                        },
-                        "return_id": True,
-                    },
+                    work_pool_name=work_pool.name,
+                    worker_name=worker.name,
+                    heartbeat_interval_seconds=30.0,
+                    get_worker_id=True,
+                    worker_metadata=WorkerMetadata(
+                        integrations=[Integration(name="prefect-aws", version="1.0.0")],
+                        custom_field="heya",
+                    ),
                 )
 
             assert worker._worker_metadata_sent