Fix tests

7a452c33510ea9db16f988daa37255b3b8b11646 · 2025-07-23 18:25:50 +0100 · Matt Beton
Files touched

D networking/target/rust-analyzer/metadata/sysroot/Cargo.lock
M shared/tests/test_sqlite_connector.py
M shared/types/events/_common.py
M shared/types/events/_events.py
M shared/types/events/chunks.py
M shared/types/events/commands.py
M shared/types/tasks.py
M worker/runner/runner_supervisor.py
M worker/tests/test_supervisor.py
M worker/tests/test_worker_handlers.py
Diff

commit 7a452c33510ea9db16f988daa37255b3b8b11646
Author: Matt Beton <matthew.beton@gmail.com>
Date:   Wed Jul 23 18:25:50 2025 +0100

    Fix tests
---
 .../rust-analyzer/metadata/sysroot/Cargo.lock      | 503 ---------------------
 shared/tests/test_sqlite_connector.py              |   2 +-
 shared/types/events/_common.py                     |   6 +
 shared/types/events/_events.py                     |   7 +-
 shared/types/events/chunks.py                      |   2 +-
 shared/types/events/commands.py                    |  16 +-
 shared/types/tasks.py                              |   4 +-
 worker/runner/runner_supervisor.py                 |   2 +-
 worker/tests/test_supervisor.py                    |  26 +-
 worker/tests/test_worker_handlers.py               |  14 +-
 10 files changed, 44 insertions(+), 538 deletions(-)

diff --git a/networking/target/rust-analyzer/metadata/sysroot/Cargo.lock b/networking/target/rust-analyzer/metadata/sysroot/Cargo.lock
deleted file mode 100644
index 97996d5f..00000000
--- a/networking/target/rust-analyzer/metadata/sysroot/Cargo.lock
+++ /dev/null
@@ -1,503 +0,0 @@
-# This file is automatically @generated by Cargo.
-# It is not intended for manual editing.
-version = 4
-
-[[package]]
-name = "addr2line"
-version = "0.22.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6e4503c46a5c0c7844e948c9a4d6acd9f50cccb4de1c48eb9e291ea17470c678"
-dependencies = [
- "compiler_builtins",
- "gimli 0.29.0",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "adler"
-version = "1.0.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "alloc"
-version = "0.0.0"
-dependencies = [
- "compiler_builtins",
- "core",
- "rand",
- "rand_xorshift",
-]
-
-[[package]]
-name = "allocator-api2"
-version = "0.2.18"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f"
-
-[[package]]
-name = "cc"
-version = "1.1.22"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9540e661f81799159abee814118cc139a2004b3a3aa3ea37724a1b66530b90e0"
-dependencies = [
- "shlex",
-]
-
-[[package]]
-name = "cfg-if"
-version = "1.0.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "compiler_builtins"
-version = "0.1.138"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "53f0ea7fff95b51f84371588f06062557e96bbe363d2b36218ddb806f3ca8611"
-dependencies = [
- "cc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "core"
-version = "0.0.0"
-dependencies = [
- "rand",
- "rand_xorshift",
-]
-
-[[package]]
-name = "dlmalloc"
-version = "0.2.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d9b5e0d321d61de16390ed273b647ce51605b575916d3c25e6ddf27a1e140035"
-dependencies = [
- "cfg-if",
- "compiler_builtins",
- "libc",
- "rustc-std-workspace-core",
- "windows-sys",
-]
-
-[[package]]
-name = "fortanix-sgx-abi"
-version = "0.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "57cafc2274c10fab234f176b25903ce17e690fca7597090d50880e047a0389c5"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "getopts"
-version = "0.2.21"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "14dbbfd5c71d70241ecf9e6f13737f7b5ce823821063188d7e46c41d371eebd5"
-dependencies = [
- "rustc-std-workspace-core",
- "rustc-std-workspace-std",
- "unicode-width",
-]
-
-[[package]]
-name = "gimli"
-version = "0.29.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "40ecd4077b5ae9fd2e9e169b102c6c330d0605168eb0e8bf79952b256dbefffd"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "gimli"
-version = "0.31.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "hashbrown"
-version = "0.15.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb"
-dependencies = [
- "allocator-api2",
- "compiler_builtins",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "hermit-abi"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "libc"
-version = "0.2.162"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "18d287de67fe55fd7e1581fe933d965a5a9477b38e949cfa9f8574ef01506398"
-dependencies = [
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "memchr"
-version = "2.7.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "miniz_oxide"
-version = "0.7.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b8a240ddb74feaf34a79a7add65a741f3167852fba007066dcac1ca548d89c08"
-dependencies = [
- "adler",
- "compiler_builtins",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "object"
-version = "0.36.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e"
-dependencies = [
- "compiler_builtins",
- "memchr",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "panic_abort"
-version = "0.0.0"
-dependencies = [
- "alloc",
- "cfg-if",
- "compiler_builtins",
- "core",
- "libc",
-]
-
-[[package]]
-name = "panic_unwind"
-version = "0.0.0"
-dependencies = [
- "alloc",
- "cfg-if",
- "compiler_builtins",
- "core",
- "libc",
- "unwind",
-]
-
-[[package]]
-name = "proc_macro"
-version = "0.0.0"
-dependencies = [
- "core",
- "std",
-]
-
-[[package]]
-name = "profiler_builtins"
-version = "0.0.0"
-dependencies = [
- "cc",
- "compiler_builtins",
- "core",
-]
-
-[[package]]
-name = "r-efi"
-version = "4.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e9e935efc5854715dfc0a4c9ef18dc69dee0ec3bf9cc3ab740db831c0fdd86a3"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "r-efi-alloc"
-version = "1.0.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "31d6f09fe2b6ad044bc3d2c34ce4979796581afd2f1ebc185837e02421e02fd7"
-dependencies = [
- "compiler_builtins",
- "r-efi",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "rand"
-version = "0.8.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
-dependencies = [
- "rand_core",
-]
-
-[[package]]
-name = "rand_core"
-version = "0.6.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
-
-[[package]]
-name = "rand_xorshift"
-version = "0.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d25bf25ec5ae4a3f1b92f929810509a2f53d7dca2f50b794ff57e3face536c8f"
-dependencies = [
- "rand_core",
-]
-
-[[package]]
-name = "rustc-demangle"
-version = "0.1.24"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "rustc-std-workspace-alloc"
-version = "1.99.0"
-dependencies = [
- "alloc",
-]
-
-[[package]]
-name = "rustc-std-workspace-core"
-version = "1.99.0"
-dependencies = [
- "core",
-]
-
-[[package]]
-name = "rustc-std-workspace-std"
-version = "1.99.0"
-dependencies = [
- "std",
-]
-
-[[package]]
-name = "shlex"
-version = "1.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
-
-[[package]]
-name = "std"
-version = "0.0.0"
-dependencies = [
- "addr2line",
- "alloc",
- "cfg-if",
- "compiler_builtins",
- "core",
- "dlmalloc",
- "fortanix-sgx-abi",
- "hashbrown",
- "hermit-abi",
- "libc",
- "miniz_oxide",
- "object",
- "panic_abort",
- "panic_unwind",
- "r-efi",
- "r-efi-alloc",
- "rand",
- "rand_xorshift",
- "rustc-demangle",
- "std_detect",
- "unwind",
- "wasi",
- "windows-targets 0.0.0",
-]
-
-[[package]]
-name = "std_detect"
-version = "0.1.5"
-dependencies = [
- "cfg-if",
- "compiler_builtins",
- "libc",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "sysroot"
-version = "0.0.0"
-dependencies = [
- "proc_macro",
- "profiler_builtins",
- "std",
- "test",
-]
-
-[[package]]
-name = "test"
-version = "0.0.0"
-dependencies = [
- "core",
- "getopts",
- "libc",
- "std",
-]
-
-[[package]]
-name = "unicode-width"
-version = "0.1.14"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-core",
- "rustc-std-workspace-std",
-]
-
-[[package]]
-name = "unwind"
-version = "0.0.0"
-dependencies = [
- "cfg-if",
- "compiler_builtins",
- "core",
- "libc",
- "unwinding",
-]
-
-[[package]]
-name = "unwinding"
-version = "0.2.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "637d511437df708cee34bdec7ba2f1548d256b7acf3ff20e0a1c559f9bf3a987"
-dependencies = [
- "compiler_builtins",
- "gimli 0.31.1",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "wasi"
-version = "0.11.0+wasi-snapshot-preview1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
-dependencies = [
- "compiler_builtins",
- "rustc-std-workspace-alloc",
- "rustc-std-workspace-core",
-]
-
-[[package]]
-name = "windows-sys"
-version = "0.59.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
-dependencies = [
- "windows-targets 0.52.6",
-]
-
-[[package]]
-name = "windows-targets"
-version = "0.0.0"
-
-[[package]]
-name = "windows-targets"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
-dependencies = [
- "windows_aarch64_gnullvm",
- "windows_aarch64_msvc",
- "windows_i686_gnu",
- "windows_i686_gnullvm",
- "windows_i686_msvc",
- "windows_x86_64_gnu",
- "windows_x86_64_gnullvm",
- "windows_x86_64_msvc",
-]
-
-[[package]]
-name = "windows_aarch64_gnullvm"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
-
-[[package]]
-name = "windows_aarch64_msvc"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
-
-[[package]]
-name = "windows_i686_gnu"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
-
-[[package]]
-name = "windows_i686_gnullvm"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
-
-[[package]]
-name = "windows_i686_msvc"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
-
-[[package]]
-name = "windows_x86_64_gnu"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
-
-[[package]]
-name = "windows_x86_64_gnullvm"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
-
-[[package]]
-name = "windows_x86_64_msvc"
-version = "0.52.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
diff --git a/shared/tests/test_sqlite_connector.py b/shared/tests/test_sqlite_connector.py
index 80736bfd..9e4c8b4d 100644
--- a/shared/tests/test_sqlite_connector.py
+++ b/shared/tests/test_sqlite_connector.py
@@ -13,10 +13,10 @@ from shared.db.sqlite import AsyncSQLiteEventStorage, EventLogConfig
 from shared.types.common import NodeId
 from shared.types.events import (
     ChunkGenerated,
+    CommandId,
     _EventType,
 )
 from shared.types.events.chunks import ChunkType, TokenChunk
-from shared.types.events.commands import CommandId
 
 # Type ignore comment for all protected member access in this test file
 # pyright: reportPrivateUsage=false
diff --git a/shared/types/events/_common.py b/shared/types/events/_common.py
index 72788da1..a5a1b18a 100644
--- a/shared/types/events/_common.py
+++ b/shared/types/events/_common.py
@@ -15,6 +15,12 @@ class EventId(NewUUID):
     """
 
 
+class CommandId(NewUUID):
+    """
+    Newtype around `NewUUID` for command IDs
+    """
+
+
 # Event base-class boilerplate (you should basically never touch these)
 # Only very specialised registry or serialisation/deserialization logic might need know about these
 
diff --git a/shared/types/events/_events.py b/shared/types/events/_events.py
index 64cefe50..07da96b9 100644
--- a/shared/types/events/_events.py
+++ b/shared/types/events/_events.py
@@ -3,13 +3,16 @@ from typing import Literal
 from shared.topology import Connection, ConnectionProfile, Node, NodePerformanceProfile
 from shared.types.common import NodeId
 from shared.types.events.chunks import GenerationChunk
-from shared.types.events.commands import CommandId
 from shared.types.tasks import Task, TaskId, TaskStatus
 from shared.types.worker.common import InstanceId, NodeStatus
 from shared.types.worker.instances import InstanceParams, TypeOfInstance
 from shared.types.worker.runners import RunnerId, RunnerStatus
 
-from ._common import _BaseEvent, _EventType  # pyright: ignore[reportPrivateUsage]
+from ._common import (
+    CommandId,
+    _BaseEvent,  # pyright: ignore[reportPrivateUsage]
+    _EventType,  # pyright: ignore[reportPrivateUsage]
+)
 
 
 class TaskCreated(_BaseEvent[_EventType.TaskCreated]):
diff --git a/shared/types/events/chunks.py b/shared/types/events/chunks.py
index 81d2bfae..e2cb7a7b 100644
--- a/shared/types/events/chunks.py
+++ b/shared/types/events/chunks.py
@@ -4,7 +4,7 @@ from typing import Annotated, Literal
 from pydantic import BaseModel, Field, TypeAdapter
 
 from shared.openai_compat import FinishReason
-from shared.types.events.commands import CommandId
+from shared.types.events._common import CommandId
 from shared.types.models import ModelId
 
 
diff --git a/shared/types/events/commands.py b/shared/types/events/commands.py
index fe645869..cce1b043 100644
--- a/shared/types/events/commands.py
+++ b/shared/types/events/commands.py
@@ -1,18 +1,15 @@
 from enum import Enum
-from typing import Annotated, Callable, Sequence
+from typing import Annotated, Callable, Literal, Sequence
 
 from pydantic import BaseModel, Field, TypeAdapter
 
 from shared.types.api import ChatCompletionTaskParams
-from shared.types.common import NewUUID
 from shared.types.events import Event
+from shared.types.events._common import CommandId
 from shared.types.state import InstanceId, State
 
 
-class CommandId(NewUUID):
-    pass
-
-
+# TODO: We need to have a distinction between create instance and spin up instance.
 class CommandTypes(str, Enum):
     CHAT_COMPLETION = "CHAT_COMPLETION"
     CREATE_INSTANCE = "CREATE_INSTANCE"
@@ -25,19 +22,22 @@ class _BaseCommand[T: CommandTypes](BaseModel):
 
 
 class ChatCompletionCommand(_BaseCommand[CommandTypes.CHAT_COMPLETION]):
+    command_type: Literal[CommandTypes.CHAT_COMPLETION] = CommandTypes.CHAT_COMPLETION
     request_params: ChatCompletionTaskParams
 
 
 class CreateInstanceCommand(_BaseCommand[CommandTypes.CREATE_INSTANCE]):
+    command_type: Literal[CommandTypes.CREATE_INSTANCE] = CommandTypes.CREATE_INSTANCE
     model_id: str
 
 
 class DeleteInstanceCommand(_BaseCommand[CommandTypes.DELETE_INSTANCE]):
+    command_type: Literal[CommandTypes.DELETE_INSTANCE] = CommandTypes.DELETE_INSTANCE
     instance_id: InstanceId
 
-
 Command = Annotated[
-    ChatCompletionCommand, Field(discriminator="command_type")
+    ChatCompletionCommand | CreateInstanceCommand | DeleteInstanceCommand, 
+    Field(discriminator="command_type")
 ]
 
 CommandParser: TypeAdapter[Command] = TypeAdapter(Command)
diff --git a/shared/types/tasks.py b/shared/types/tasks.py
index 011f084a..08e9e017 100644
--- a/shared/types/tasks.py
+++ b/shared/types/tasks.py
@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Annotated
+from typing import Annotated, Literal
 
 from pydantic import BaseModel, Field
 
@@ -24,7 +24,7 @@ class TaskStatus(str, Enum):
 
 
 class ChatCompletionTask(BaseModel):
-    task_type: TaskType
+    task_type: Literal[TaskType.CHAT_COMPLETION] = TaskType.CHAT_COMPLETION
     task_id: TaskId
     instance_id: InstanceId
     task_status: TaskStatus
diff --git a/worker/runner/runner_supervisor.py b/worker/runner/runner_supervisor.py
index 1f60f1d9..7e69358f 100644
--- a/worker/runner/runner_supervisor.py
+++ b/worker/runner/runner_supervisor.py
@@ -5,8 +5,8 @@ from collections.abc import AsyncGenerator
 from types import CoroutineType
 from typing import Any, Callable
 
+from shared.types.events import CommandId
 from shared.types.events.chunks import GenerationChunk, TokenChunk
-from shared.types.events.commands import CommandId
 from shared.types.tasks import ChatCompletionTaskParams, Task
 from shared.types.worker.commands_runner import (
     ChatTaskMessage,
diff --git a/worker/tests/test_supervisor.py b/worker/tests/test_supervisor.py
index 40f4ba02..b482e833 100644
--- a/worker/tests/test_supervisor.py
+++ b/worker/tests/test_supervisor.py
@@ -26,7 +26,7 @@ def user_message():
 async def test_supervisor_single_node_response(
     pipeline_shard_meta: Callable[..., PipelineShardMetadata],
     hosts: Callable[..., list[Host]],
-    chat_task: Task,
+    chat_completion_task: Task,
     tmp_path: Path,
 ):
     """Test that asking for the capital of France returns 'Paris' in the response"""
@@ -43,7 +43,7 @@ async def test_supervisor_single_node_response(
         full_response = ""
         stop_reason: FinishReason | None = None
 
-        async for chunk in supervisor.stream_response(task=chat_task):
+        async for chunk in supervisor.stream_response(task=chat_completion_task):
             if isinstance(chunk, TokenChunk):
                 full_response += chunk.text
                 if chunk.finish_reason:
@@ -63,7 +63,7 @@ async def test_supervisor_single_node_response(
 async def test_supervisor_two_node_response(
     pipeline_shard_meta: Callable[..., PipelineShardMetadata],
     hosts: Callable[..., list[Host]],
-    chat_task: Task,
+    chat_completion_task: Task,
     tmp_path: Path,
 ):
     """Test that asking for the capital of France returns 'Paris' in the response"""
@@ -85,13 +85,13 @@ async def test_supervisor_two_node_response(
 
         async def collect_response_0():
             nonlocal full_response_0
-            async for chunk in supervisor_0.stream_response(task=chat_task):
+            async for chunk in supervisor_0.stream_response(task=chat_completion_task):
                 if isinstance(chunk, TokenChunk):
                     full_response_0 += chunk.text
 
         async def collect_response_1():
             nonlocal full_response_1
-            async for chunk in supervisor_1.stream_response(task=chat_task):
+            async for chunk in supervisor_1.stream_response(task=chat_completion_task):
                 if isinstance(chunk, TokenChunk):
                     full_response_1 += chunk.text
 
@@ -118,7 +118,7 @@ async def test_supervisor_two_node_response(
 async def test_supervisor_early_stopping(
     pipeline_shard_meta: Callable[..., PipelineShardMetadata],
     hosts: Callable[..., list[Host]],
-    chat_task: Task,
+    chat_completion_task: Task,
     tmp_path: Path,
 ):
     """Test that asking for the capital of France returns 'Paris' in the response"""
@@ -130,10 +130,10 @@ async def test_supervisor_early_stopping(
     )
 
     max_tokens = 50
-    assert chat_task.task_type == TaskType.CHAT_COMPLETION
-    print(f'chat_task.task_params: {chat_task.task_params}')
-    assert isinstance(chat_task.task_params, ChatCompletionTaskParams)
-    task_params: ChatCompletionTaskParams = chat_task.task_params
+    assert chat_completion_task.task_type == TaskType.CHAT_COMPLETION
+    print(f'chat_completion_task.task_params: {chat_completion_task.task_params}')
+    assert isinstance(chat_completion_task.task_params, ChatCompletionTaskParams)
+    task_params: ChatCompletionTaskParams = chat_completion_task.task_params
 
     try:
         task_params.max_tokens = max_tokens
@@ -146,7 +146,7 @@ async def test_supervisor_early_stopping(
         count = 0
         stop_reason: FinishReason | None = None
 
-        async for chunk in supervisor.stream_response(task=chat_task):
+        async for chunk in supervisor.stream_response(task=chat_completion_task):
             if isinstance(chunk, TokenChunk):
                 full_response += chunk.text
                 count += 1
@@ -169,7 +169,7 @@ async def test_supervisor_early_stopping(
 async def test_supervisor_handles_terminated_runner(
     pipeline_shard_meta: Callable[..., PipelineShardMetadata],
     hosts: Callable[..., list[Host]],
-    chat_task: Task,
+    chat_completion_task: Task,
     tmp_path: Path,
 ):
     """Test that the supervisor handles a terminated runner"""
@@ -194,7 +194,7 @@ async def test_supervisor_handles_terminated_runner(
 async def test_supervisor_handles_killed_runner(
     pipeline_shard_meta: Callable[..., PipelineShardMetadata],
     hosts: Callable[..., list[Host]],
-    chat_task: Task,
+    chat_completion_task: Task,
     tmp_path: Path,
 ):
     """Test that the supervisor handles a killed runner"""
diff --git a/worker/tests/test_worker_handlers.py b/worker/tests/test_worker_handlers.py
index e1a01ca3..20823c5e 100644
--- a/worker/tests/test_worker_handlers.py
+++ b/worker/tests/test_worker_handlers.py
@@ -84,7 +84,7 @@ async def test_unassign_op(worker_with_assigned_runner: tuple[Worker, RunnerId,
     assert len(events) == 0
 
 @pytest.mark.asyncio
-async def test_runner_up_op(worker_with_assigned_runner: tuple[Worker, RunnerId, Instance], chat_task: Task, tmp_path: Path):
+async def test_runner_up_op(worker_with_assigned_runner: tuple[Worker, RunnerId, Instance], chat_completion_task: Task, tmp_path: Path):
     worker, runner_id, _ = worker_with_assigned_runner
 
     runner_up_op = RunnerUpOp(runner_id=runner_id)
@@ -104,7 +104,7 @@ async def test_runner_up_op(worker_with_assigned_runner: tuple[Worker, RunnerId,
 
     full_response = ''
 
-    async for chunk in supervisor.stream_response(task=chat_task):
+    async for chunk in supervisor.stream_response(task=chat_completion_task):
         if isinstance(chunk, TokenChunk):
             full_response += chunk.text
 
@@ -153,12 +153,12 @@ async def test_download_op(worker_with_assigned_runner: tuple[Worker, RunnerId,
 @pytest.mark.asyncio
 async def test_execute_task_op(
     worker_with_running_runner: tuple[Worker, RunnerId, Instance], 
-    chat_task: Task, tmp_path: Path):
+    chat_completion_task: Task, tmp_path: Path):
     worker, runner_id, _ = worker_with_running_runner
 
     execute_task_op = ExecuteTaskOp(
         runner_id=runner_id,
-        task=chat_task
+        task=chat_completion_task
     )
 
     events: list[Event] = []
@@ -187,15 +187,15 @@ async def test_execute_task_op(
 @pytest.mark.asyncio
 async def test_execute_task_fails(
     worker_with_running_runner: tuple[Worker, RunnerId, Instance], 
-    chat_task: Task, tmp_path: Path):
+    chat_completion_task: Task, tmp_path: Path):
     worker, runner_id, _ = worker_with_running_runner
 
-    messages = chat_task.task_params.messages
+    messages = chat_completion_task.task_params.messages
     messages[0].content = 'Artificial prompt: EXO RUNNER MUST FAIL'
 
     execute_task_op = ExecuteTaskOp(
         runner_id=runner_id,
-        task=chat_task
+        task=chat_completion_task
     )
 
     events: list[Event] = []
← 7ac23ce9 Refactor tasks / commands / api · back to Exo · wrote race-condition-free persistent NodeID-getting function 3ab56092 →