fix: wait for TASK_STATE_RUNNING before exec_in_container

The test assumed task state matches job state, but tasks can be
BUILDING while the job reports RUNNING (workspace bundle install).
Add wait_for_task_state helper and use it.

Author: rjpower

tests/integration/iris/cluster.py

@@ -107,6 +107,25 @@ def wait_for_state(
             time.sleep(poll_interval)
         raise TimeoutError(f"Job {job.job_id} did not reach state {state} in {timeout}s (current: {status.state})")
 
+    def wait_for_task_state(
+        self,
+        job: Job,
+        state: int,
+        task_index: int = 0,
+        timeout: float = 60.0,
+        poll_interval: float = 0.5,
+    ) -> cluster_pb2.TaskStatus:
+        deadline = time.monotonic() + timeout
+        task = self.task_status(job, task_index)
+        while time.monotonic() < deadline:
+            task = self.task_status(job, task_index)
+            if task.state == state:
+                return task
+            time.sleep(poll_interval)
+        raise TimeoutError(
+            f"Task {task_index} of {job.job_id} did not reach state {state} " f"in {timeout}s (current: {task.state})"
+        )
+
     @contextmanager
     def launched_job(self, fn, name: str, *args, **kwargs):
         """Submit a job and guarantee it's killed on exit."""

tests/integration/iris/test_iris_integration.py

@@ -258,10 +258,9 @@ def _is_running():
 def test_exec_in_container(integration_cluster):
     """Exec a command in a running task's container."""
     job = integration_cluster.submit(sleep, "itest-exec", 120)
-    integration_cluster.wait_for_state(job, cluster_pb2.JOB_STATE_RUNNING, timeout=integration_cluster.job_timeout)
-
-    task = integration_cluster.task_status(job, task_index=0)
-    assert task.state == cluster_pb2.TASK_STATE_RUNNING, f"Task stuck in {cluster_pb2.TaskState.Name(task.state)}"
+    task = integration_cluster.wait_for_task_state(
+        job, cluster_pb2.TASK_STATE_RUNNING, timeout=integration_cluster.job_timeout
+    )
     task_id = task.task_id
 
     request = cluster_pb2.Controller.ExecInContainerRequest(