continue implementation

testarena_app/main.py

@@ -31,10 +31,13 @@ models.Base.metadata.create_all(bind=database.engine)
 @app.post("/api/queue")
 async def queue_task(payload: Dict, db: Session = Depends(database.get_db)):
     """
-    Input json contain {<queue_ID> :[environment, "<TASK_ID>" : "<path to scenario>],}
+    Input json contain {"source": "<branch_name>", <queue_ID> :[environment, {"<TASK_ID>" : "<path to scenario>"},]}
     """
     try:
-        queue_id = list(payload.keys())[0]
+        source = payload.get("source", "main")
+        # Find the queue_id key (it's the one that isn't "source")
+        queue_id = next(k for k in payload.keys() if k != "source")
+        
         data = payload[queue_id]
         environment = data[0]
         tasks_data = data[1] # This is a dict {"TASK_ID": "path"}
@@ -47,12 +50,13 @@ async def queue_task(payload: Dict, db: Session = Depends(database.get_db)):
         status_file = os.path.join(queue_dir, "queue_status.json")
         queue_status = {
             "queue_id": queue_id,
+            "source": source,
             "status": "Waiting",
             "tasks": {}
         }
         
         # 3. Save to database and prepare status file
-        new_queue = models.Queue(id=queue_id, environment=environment, status="Waiting")
+        new_queue = models.Queue(id=queue_id, environment=environment, source=source, status="Waiting")
         db.add(new_queue)
         
         for task_id, scenario_path in tasks_data.items():