test: tighten scheduler queue accounting coverage

Signed-off-by: Eric W. Tramel <1223539+eric-tramel@users.noreply.github.com>

Author: eric-tramel

packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py

@@ -145,9 +145,7 @@ def view(self) -> QueueView:
         first_tasks_by_group: dict[TaskGroupKey, SchedulableTask] = {}
         first_group_specs: dict[TaskGroupKey, TaskGroupSpec] = {}
 
-        for key, queue in self._queues.items():
-            if self._queued_by_group.get(key, 0) <= 0:
-                continue
+        for key in self._queued_by_group:
             first = self._first_valid_item(key)
             if first is None:
                 continue
@@ -157,18 +155,14 @@ def view(self) -> QueueView:
 
         return QueueView(
             queued_total=len(self._queued),
-            queued_by_group={key: count for key, count in self._queued_by_group.items() if count > 0},
+            queued_by_group=dict(self._queued_by_group),
             queued_resource_demand_by_group={
-                key: {resource: count for resource, count in value.items() if count > 0}
-                for key, value in self._queued_resource_demand_by_group.items()
-                if self._queued_by_group.get(key, 0) > 0
+                key: dict(value) for key, value in self._queued_resource_demand_by_group.items()
             },
             first_candidate_resources_by_group=first_by_group,
             first_candidate_tasks_by_group=first_tasks_by_group,
             first_candidate_group_specs_by_group=first_group_specs,
-            queued_peer_demand_by_resource={
-                resource: count for resource, count in self._queued_peer_demand_by_resource.items() if count > 0
-            },
+            queued_peer_demand_by_resource=dict(self._queued_peer_demand_by_resource),
         )
 
     def _activate_group(self, key: TaskGroupKey) -> None:

packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_queue.py

@@ -129,14 +129,25 @@ def test_select_next_uses_scheduler_eligibility_callback() -> None:
 
 def test_enqueue_is_idempotent_by_task_id() -> None:
     queue = FairTaskQueue()
-    item = _item("a", 0)
+    group = _group("a")
+    task = _task("a", 0)
+    item = SchedulableTask(
+        task_id=stable_task_id(task),
+        payload=task,
+        group=group,
+        resource_request=SchedulerResourceRequest({"submission": 1, "llm_wait": 2}),
+    )
 
     first = queue.enqueue([item])
     second = queue.enqueue([item])
+    view = queue.view()
 
     assert first == (item.task_id,)
     assert second == ()
-    assert queue.view().queued_total == 1
+    assert view.queued_total == 1
+    assert view.queued_by_group == {group.key: 1}
+    assert view.queued_resource_demand_by_group == {group.key: {"submission": 1, "llm_wait": 2}}
+    assert view.queued_peer_demand_by_resource == {"submission": 1, "llm_wait": 2}
 
 
 def test_discard_where_removes_matching_tasks() -> None:
@@ -208,28 +219,23 @@ def test_queue_view_updates_incremental_accounting_after_removals() -> None:
 
 def test_queue_view_uses_incremental_accounting_for_non_candidate_tasks() -> None:
     queue = FairTaskQueue()
-    guarded_amounts: list[_FailIfScannedAmounts] = []
-    items: list[SchedulableTask] = []
-    for group_index in range(8):
-        group = _group(f"group-{group_index}")
-        items.append(_item(f"group-{group_index}", 0, group))
-        for row in range(1, 32):
-            amounts = _FailIfScannedAmounts({"submission": 1})
-            task = _task(f"group-{group_index}", row)
-            items.append(
-                SchedulableTask(
-                    task_id=stable_task_id(task),
-                    payload=task,
-                    group=group,
-                    resource_request=SchedulerResourceRequest(amounts),
-                )
-            )
-            guarded_amounts.append(amounts)
-    queue.enqueue(items)
-    for amounts in guarded_amounts:
-        amounts.locked = True
+    group = _group("a")
+    first = _item("a", 0, group)
+    amounts = _FailIfScannedAmounts({"submission": 1})
+    task = _task("a", 1)
+    second = SchedulableTask(
+        task_id=stable_task_id(task),
+        payload=task,
+        group=group,
+        resource_request=SchedulerResourceRequest(amounts),
+    )
+    queue.enqueue([first, second])
+    amounts.locked = True
 
     view = queue.view()
 
-    assert view.queued_total == 256
-    assert sum(view.queued_by_group.values()) == 256
+    assert view.queued_total == 2
+    assert view.queued_by_group == {group.key: 2}
+    assert view.queued_resource_demand_by_group == {group.key: {"submission": 2}}
+    assert view.first_candidate_resources_by_group == {group.key: {"submission": 1}}
+    assert view.queued_peer_demand_by_resource == {"submission": 2}