From d2bbe4de9545f6a34c7138a3b3a2feccfa7bbf4d Mon Sep 17 00:00:00 2001
From: Laura Abro <laura@koii.network>
Date: Thu, 24 Apr 2025 11:52:26 -0300
Subject: [PATCH] add worker fetch step

---
 .../orca-agent/tests/stages/worker_fetch.py   | 34 ++++---------------
 worker/orca-agent/tests/steps.py              | 10 +++++-
 2 files changed, 16 insertions(+), 28 deletions(-)

diff --git a/worker/orca-agent/tests/stages/worker_fetch.py b/worker/orca-agent/tests/stages/worker_fetch.py
index 294f458..fc625a2 100644
--- a/worker/orca-agent/tests/stages/worker_fetch.py
+++ b/worker/orca-agent/tests/stages/worker_fetch.py
@@ -28,7 +28,7 @@ def prepare(runner, worker):
 
 def execute(runner, worker, data):
     """Execute worker task step"""
-    url = f"{worker.url}/worker-task/{data['roundNumber']}"
+    url = f"{runner.config.middle_server_url}/summarizer/worker/fetch-todo"
     response = requests.post(
         url,
         json={"signature": data["stakingSignature"], "stakingKey": data["stakingKey"]},
@@ -41,32 +41,12 @@ def execute(runner, worker, data):
             f"✓ {result.get('message', 'No eligible todos')} for {worker.name} - continuing"
         )
         return {"success": True, "message": result.get("message")}
+    else:
+        response.raise_for_status()
 
-    if result.get("success") and "pr_url" in result:
-        round_key = str(runner.current_round)
-        round_state = runner.state["rounds"].setdefault(round_key, {})
-
-        # Initialize pr_urls if not exists
-        if "pr_urls" not in round_state:
-            round_state["pr_urls"] = {}
-        round_state["pr_urls"][worker.name] = result["pr_url"]
-
-        # Initialize submission_data if not exists
-        if "submission_data" not in round_state:
-            round_state["submission_data"] = {}
-
-        # Store submission data
-        round_state["submission_data"][worker.name] = {
-            "githubUsername": worker.env.get("GITHUB_USERNAME"),
-            "nodeType": "worker",
-            "prUrl": result["pr_url"],
-            "repoName": result.get("repoName"),
-            "repoOwner": result.get("repoOwner"),
-            "roundNumber": runner.current_round,
-            "taskId": runner.config.task_id,
-            "uuid": result.get("uuid"),  # Should be provided by the worker
-            "stakingKey": worker.staking_public_key,
-            "pubKey": worker.public_key,
-        }
+    if result.get("success"):
+        runner.state["repo_url"] = (
+            f"https://github.com/{result['repo_owner']}/{result['repo_name']}"
+        )
 
     return result
diff --git a/worker/orca-agent/tests/steps.py b/worker/orca-agent/tests/steps.py
index 3af9843..e669361 100644
--- a/worker/orca-agent/tests/steps.py
+++ b/worker/orca-agent/tests/steps.py
@@ -3,6 +3,7 @@
 from prometheus_test import TestStep
 from functools import partial
 from .stages import (
+    worker_fetch,
     worker_task,
     worker_submission,
     worker_audit,
@@ -11,12 +12,19 @@ from .stages import (
 
 
 steps = [
+    TestStep(
+        name="worker_fetch",
+        description="Fetch worker task",
+        prepare=worker_fetch.prepare,
+        execute=worker_fetch.execute,
+        worker="worker1",
+    ),
     TestStep(
         name="worker_task",
         description="Execute worker task",
         prepare=worker_task.prepare,
         execute=worker_task.execute,
-        worker="worker",
+        worker="worker1",
     ),
     TestStep(
         name="worker_submission",