Correctly display job creation errors (#133)

* load error response json * fix tests * face the consequences of copy_model_fields banishment
PrefectHQ · Apr 10, 2024 · 8cafd9f · 8cafd9f
1 parent 0a16944
commit 8cafd9f
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 42 deletions.
diff --git a/prefect_kubernetes/worker.py b/prefect_kubernetes/worker.py
@@ -103,6 +103,7 @@
 import asyncio
 import base64
 import enum
+import json
 import logging
 import math
 import os
@@ -795,8 +796,8 @@ def _create_job(
             message = ""
             if exc.reason:
                 message += ": " + exc.reason
-            if exc.body and "message" in exc.body:
-                message += ": " + exc.body["message"]
+            if exc.body and "message" in (body := json.loads(exc.body)):
+                message += ": " + body["message"]
 
             raise InfrastructureError(
                 f"Unable to create Kubernetes job{message}"

diff --git a/tests/test_worker.py b/tests/test_worker.py
@@ -1,4 +1,5 @@
 import base64
+import json
 import re
 import uuid
 from contextlib import contextmanager
@@ -1008,7 +1009,7 @@ def flow_run(self):
 
     @pytest.fixture
     def deployment(self):
-        return DeploymentResponse(name="my-deployment-name")
+        return DeploymentResponse(name="my-deployment-name", flow_id=uuid.uuid4())
 
     @pytest.fixture
     def flow(self):
@@ -1522,16 +1523,18 @@ async def test_create_job_failure(
         mock_batch_client,
     ):
         response = MagicMock()
-        response.data = {
-            "kind": "Status",
-            "apiVersion": "v1",
-            "metadata": {},
-            "status": "Failure",
-            "message": 'jobs.batch is forbidden: User "system:serviceaccount:helm-test:prefect-worker-dev" cannot create resource "jobs" in API group "batch" in the namespace "prefect"',
-            "reason": "Forbidden",
-            "details": {"group": "batch", "kind": "jobs"},
-            "code": 403,
-        }
+        response.data = json.dumps(
+            {
+                "kind": "Status",
+                "apiVersion": "v1",
+                "metadata": {},
+                "status": "Failure",
+                "message": 'jobs.batch is forbidden: User "system:serviceaccount:helm-test:prefect-worker-dev" cannot create resource "jobs" in API group "batch" in the namespace "prefect"',
+                "reason": "Forbidden",
+                "details": {"group": "batch", "kind": "jobs"},
+                "code": 403,
+            }
+        )
         response.status = 403
         response.reason = "Forbidden"
 
@@ -1563,16 +1566,18 @@ async def test_create_job_retries(
     ):
         MAX_ATTEMPTS = 3
         response = MagicMock()
-        response.data = {
-            "kind": "Status",
-            "apiVersion": "v1",
-            "metadata": {},
-            "status": "Failure",
-            "message": 'jobs.batch is forbidden: User "system:serviceaccount:helm-test:prefect-worker-dev" cannot create resource "jobs" in API group "batch" in the namespace "prefect"',
-            "reason": "Forbidden",
-            "details": {"group": "batch", "kind": "jobs"},
-            "code": 403,
-        }
+        response.data = json.dumps(
+            {
+                "kind": "Status",
+                "apiVersion": "v1",
+                "metadata": {},
+                "status": "Failure",
+                "message": 'jobs.batch is forbidden: User "system:serviceaccount:helm-test:prefect-worker-dev" cannot create resource "jobs" in API group "batch" in the namespace "prefect"',
+                "reason": "Forbidden",
+                "details": {"group": "batch", "kind": "jobs"},
+                "code": 403,
+            }
+        )
         response.status = 403
         response.reason = "Forbidden"
 
@@ -1605,16 +1610,18 @@ async def test_create_job_failure_no_reason(
         mock_batch_client,
     ):
         response = MagicMock()
-        response.data = {
-            "kind": "Status",
-            "apiVersion": "v1",
-            "metadata": {},
-            "status": "Failure",
-            "message": 'jobs.batch is forbidden: User "system:serviceaccount:helm-test:prefect-worker-dev" cannot create resource "jobs" in API group "batch" in the namespace "prefect"',
-            "reason": "Forbidden",
-            "details": {"group": "batch", "kind": "jobs"},
-            "code": 403,
-        }
+        response.data = json.dumps(
+            {
+                "kind": "Status",
+                "apiVersion": "v1",
+                "metadata": {},
+                "status": "Failure",
+                "message": 'jobs.batch is forbidden: User "system:serviceaccount:helm-test:prefect-worker-dev" cannot create resource "jobs" in API group "batch" in the namespace "prefect"',
+                "reason": "Forbidden",
+                "details": {"group": "batch", "kind": "jobs"},
+                "code": 403,
+            }
+        )
         response.status = 403
         response.reason = None
 
@@ -1645,15 +1652,17 @@ async def test_create_job_failure_no_message(
         mock_batch_client,
     ):
         response = MagicMock()
-        response.data = {
-            "kind": "Status",
-            "apiVersion": "v1",
-            "metadata": {},
-            "status": "Failure",
-            "reason": "Forbidden",
-            "details": {"group": "batch", "kind": "jobs"},
-            "code": 403,
-        }
+        response.data = json.dumps(
+            {
+                "kind": "Status",
+                "apiVersion": "v1",
+                "metadata": {},
+                "status": "Failure",
+                "reason": "Forbidden",
+                "details": {"group": "batch", "kind": "jobs"},
+                "code": 403,
+            }
+        )
         response.status = 403
         response.reason = "Test"