From 5c41ec91190fea60a55ad83cf11ac9e11df760d6 Mon Sep 17 00:00:00 2001
From: Laura Abro <laura@koii.network>
Date: Thu, 24 Apr 2025 22:00:54 -0300
Subject: [PATCH] fix edge cases when data doesn't exist

---
 worker/orca-agent/tests/stages/update_audit.py |  3 ++-
 worker/orca-agent/tests/stages/worker_check.py |  9 +++++++--
 worker/orca-agent/tests/stages/worker_fetch.py |  3 ++-
 worker/orca-agent/tests/stages/worker_pr.py    |  4 ++--
 worker/orca-agent/tests/stages/worker_task.py  | 11 ++++++-----
 5 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/worker/orca-agent/tests/stages/update_audit.py b/worker/orca-agent/tests/stages/update_audit.py
index c3f112b..c9fe3bd 100644
--- a/worker/orca-agent/tests/stages/update_audit.py
+++ b/worker/orca-agent/tests/stages/update_audit.py
@@ -21,4 +21,5 @@ def execute(runner, worker, data):
     )
     response.raise_for_status()
 
-    return response.json()
+    # Return a formatted response regardless of type
+    return {"success": True, "message": response.text}
diff --git a/worker/orca-agent/tests/stages/worker_check.py b/worker/orca-agent/tests/stages/worker_check.py
index fcc1d54..86fab89 100644
--- a/worker/orca-agent/tests/stages/worker_check.py
+++ b/worker/orca-agent/tests/stages/worker_check.py
@@ -6,17 +6,22 @@ import requests
 def prepare(runner, worker):
     """Prepare data for worker task"""
     # Create fetch-todo payload for stakingSignature and publicSignature
-    round_state = runner.state["rounds"][str(runner.current_round)]
+    round_state = runner.state["rounds"].get(str(runner.current_round), {})
+    if not round_state.get("pr_urls"):
+        print(f"✓ No PR URLs found for {worker.name} - continuing")
+        return
     return {
         "stakingKey": worker.staking_public_key,
         "roundNumber": runner.current_round,
         "githubUsername": worker.env.get("GITHUB_USERNAME"),
-        "prUrl": round_state["pr_urls"][worker.name],
+        "prUrl": round_state.get("pr_urls", {}).get(worker.name),
     }
 
 
 def execute(runner, worker, data):
     """Execute worker task step"""
+    if not data:
+        return {"success": True, "message": "No PR URL found"}
     url = f"{runner.config.middle_server_url}/summarizer/worker/check-todo"
     response = requests.post(
         url,
diff --git a/worker/orca-agent/tests/stages/worker_fetch.py b/worker/orca-agent/tests/stages/worker_fetch.py
index 27a4f1c..c7c72e8 100644
--- a/worker/orca-agent/tests/stages/worker_fetch.py
+++ b/worker/orca-agent/tests/stages/worker_fetch.py
@@ -45,7 +45,8 @@ def execute(runner, worker, data):
         response.raise_for_status()
 
     if result.get("success"):
-        runner.state["repo_url"] = (
+        round_state = runner.state["rounds"][runner.current_round]
+        round_state["repo_url"] = (
             f"https://github.com/{result['data']['repo_owner']}/{result['data']['repo_name']}"
         )
 
diff --git a/worker/orca-agent/tests/stages/worker_pr.py b/worker/orca-agent/tests/stages/worker_pr.py
index ff61866..50f9ca9 100644
--- a/worker/orca-agent/tests/stages/worker_pr.py
+++ b/worker/orca-agent/tests/stages/worker_pr.py
@@ -3,9 +3,9 @@ from prometheus_test.utils import create_signature
 
 
 def prepare(runner, worker):
-    round_state = runner.state["rounds"][str(runner.current_round)]
+    round_state = runner.state["rounds"].get(str(runner.current_round), {})
 
-    if worker.name not in round_state["pr_urls"]:
+    if worker.name not in round_state.get("pr_urls", {}):
         print(f"✓ No PR URL found for {worker.name} - continuing")
         return None
 
diff --git a/worker/orca-agent/tests/stages/worker_task.py b/worker/orca-agent/tests/stages/worker_task.py
index 4057703..c8c9ca6 100644
--- a/worker/orca-agent/tests/stages/worker_task.py
+++ b/worker/orca-agent/tests/stages/worker_task.py
@@ -5,20 +5,21 @@ import requests
 
 def prepare(runner, worker):
     """Prepare data for worker task"""
-
+    round_state = runner.state["rounds"].get(str(runner.current_round), {})
+    if not round_state.get("repo_url"):
+        print(f"✓ No repo url found for {worker.name} - continuing")
+        return
     return {
         "taskId": runner.config.task_id,
         "round_number": str(runner.current_round),
-        "repo_url": runner.state["repo_url"],
+        "repo_url": round_state["repo_url"],
     }
 
 
 def execute(runner, worker, data):
     """Execute worker task step"""
-    if not runner.state["repo_url"]:
-        print(f"✓ No repo url found for {worker.name} - continuing")
+    if not data:
         return {"success": True, "message": "No repo url found"}
-
     url = f"{worker.url}/worker-task/{runner.current_round}"
     response = requests.post(url, json=data)
     result = response.json()