feat(tests): operator assignment — two queues for techs, assign endpoint for leads

backend/app/routers/tests.py

@@ -51,6 +51,7 @@ from app.models.user import User
 
 # Import  from app.schemas.test
 from app.schemas.test import (
+    TestAssign,
     TestBlueUpdate,
     TestBlueValidate,
     TestClassificationUpdate,
@@ -563,6 +564,15 @@ def update_test_red(
     Returns:
         TestOut: The updated test with refreshed red-team field values.
     """
+    # Assignee lock: red_tech cannot work a test assigned to someone else
+    _pre_test = crud_get_test_or_raise(db, test_id)
+    if (
+        _pre_test.red_tech_assignee is not None
+        and current_user.role == "red_tech"
+        and _pre_test.red_tech_assignee != current_user.id
+    ):
+        raise HTTPException(status_code=403, detail="Test is assigned to another operator")
+
     # Assign update_data = payload.model_dump(exclude_unset=True)
     update_data = payload.model_dump(exclude_unset=True)
     # Open context manager
@@ -620,6 +630,15 @@ def update_test_blue(
     Returns:
         TestOut: The updated test with refreshed blue-team field values.
     """
+    # Assignee lock: blue_tech cannot work a test assigned to someone else
+    _pre_test = crud_get_test_or_raise(db, test_id)
+    if (
+        _pre_test.blue_tech_assignee is not None
+        and current_user.role == "blue_tech"
+        and _pre_test.blue_tech_assignee != current_user.id
+    ):
+        raise HTTPException(status_code=403, detail="Test is assigned to another operator")
+
     # Assign update_data = payload.model_dump(exclude_unset=True)
     update_data = payload.model_dump(exclude_unset=True)
     # Open context manager
@@ -676,10 +695,22 @@ def start_execution(
     """
     # Assign test = crud_get_test_or_raise(db, test_id)
     test = crud_get_test_or_raise(db, test_id)
+
+    # Assignee lock: red_tech cannot start a test assigned to someone else
+    if (
+        test.red_tech_assignee is not None
+        and current_user.role == "red_tech"
+        and test.red_tech_assignee != current_user.id
+    ):
+        raise HTTPException(status_code=403, detail="Test is assigned to another operator")
+
     # Open context manager
     with UnitOfWork(db) as uow:
         # Assign test = wf_start_execution(db, test, current_user)
         test = wf_start_execution(db, test, current_user)
+        # Auto-assign: if no assignee set, claim this test for the current operator
+        if test.red_tech_assignee is None:
+            test.red_tech_assignee = current_user.id
         # Call uow.commit()
         uow.commit()
     # Reload ORM object attributes from the database
@@ -715,6 +746,15 @@ def submit_red(
     """
     # Assign test = crud_get_test_or_raise(db, test_id)
     test = crud_get_test_or_raise(db, test_id)
+
+    # Assignee lock: red_tech cannot submit a test assigned to someone else
+    if (
+        test.red_tech_assignee is not None
+        and current_user.role == "red_tech"
+        and test.red_tech_assignee != current_user.id
+    ):
+        raise HTTPException(status_code=403, detail="Test is assigned to another operator")
+
     # Open context manager
     with UnitOfWork(db) as uow:
         # Assign test = wf_submit_red(db, test, current_user)
@@ -754,6 +794,15 @@ def submit_blue(
     """
     # Assign test = crud_get_test_or_raise(db, test_id)
     test = crud_get_test_or_raise(db, test_id)
+
+    # Assignee lock: blue_tech cannot submit a test assigned to someone else
+    if (
+        test.blue_tech_assignee is not None
+        and current_user.role == "blue_tech"
+        and test.blue_tech_assignee != current_user.id
+    ):
+        raise HTTPException(status_code=403, detail="Test is assigned to another operator")
+
     # Open context manager
     with UnitOfWork(db) as uow:
         # Assign test = wf_submit_blue(db, test, current_user)
@@ -779,8 +828,20 @@ def start_blue_work(
 ):
     """Blue tech picks up the test to start evaluating. Sets the Tempo timer start."""
     test = crud_get_test_or_raise(db, test_id)
+
+    # Assignee lock: blue_tech cannot pick up a test assigned to someone else
+    if (
+        test.blue_tech_assignee is not None
+        and current_user.role == "blue_tech"
+        and test.blue_tech_assignee != current_user.id
+    ):
+        raise HTTPException(status_code=403, detail="Test is assigned to another operator")
+
     with UnitOfWork(db) as uow:
         test = wf_start_blue_work(db, test, current_user)
+        # Auto-assign: if no assignee set, claim this test for the current operator
+        if test.blue_tech_assignee is None:
+            test.blue_tech_assignee = current_user.id
         uow.commit()
     db.refresh(test)
     return test
@@ -1017,6 +1078,48 @@ def reopen(
     return test
 
 
+# ---------------------------------------------------------------------------
+# POST /tests/{id}/assign  —  assign red_tech / blue_tech operators (leads + admin)
+# ---------------------------------------------------------------------------
+
+
+@router.post("/{test_id}/assign", response_model=TestOut)
+def assign_test_operators(
+    test_id: uuid.UUID,
+    payload: TestAssign,
+    db: Session = Depends(get_db),
+    current_user: User = Depends(require_any_role("admin", "red_lead", "blue_lead")),
+):
+    """Assign red_tech and/or blue_tech operators to a test. Admin/leads only."""
+    test = crud_get_test_or_raise(db, test_id)
+
+    if payload.red_tech_assignee is not None:
+        u = db.query(User).filter(User.id == payload.red_tech_assignee).first()
+        if not u or u.role not in ("red_tech", "red_lead", "admin"):
+            raise HTTPException(status_code=400, detail="Invalid red tech assignee")
+        test.red_tech_assignee = payload.red_tech_assignee
+
+    if payload.blue_tech_assignee is not None:
+        u = db.query(User).filter(User.id == payload.blue_tech_assignee).first()
+        if not u or u.role not in ("blue_tech", "blue_lead", "admin"):
+            raise HTTPException(status_code=400, detail="Invalid blue tech assignee")
+        test.blue_tech_assignee = payload.blue_tech_assignee
+
+    # Handle intentional null (clearing) — model_fields_set tracks which keys were sent
+    if "red_tech_assignee" in payload.model_fields_set and payload.red_tech_assignee is None:
+        test.red_tech_assignee = None
+    if "blue_tech_assignee" in payload.model_fields_set and payload.blue_tech_assignee is None:
+        test.blue_tech_assignee = None
+
+    log_action(db, current_user.id, "assign_test", str(test_id), {
+        "red_tech_assignee": str(payload.red_tech_assignee) if payload.red_tech_assignee else None,
+        "blue_tech_assignee": str(payload.blue_tech_assignee) if payload.blue_tech_assignee else None,
+    })
+    db.commit()
+    db.refresh(test)
+    return test
+
+
 # ---------------------------------------------------------------------------
 # PATCH /tests/{id}/remediation  —  update remediation fields
 # ---------------------------------------------------------------------------