Tests refactoring: explcit worker ids removed

Author: spirali

crates/tako/src/internal/server/core.rs

@@ -632,66 +632,60 @@ mod tests {
             self.single_node_ready_to_assign.retain(|&id| id != task_id);
         }
 
-        pub fn assert_task_condition<T: Copy + Into<TaskId>, F: Fn(&Task) -> bool>(
-            &self,
-            task_ids: &[T],
-            op: F,
-        ) {
+        pub fn assert_task_condition<F: Fn(&Task) -> bool>(&self, task_ids: &[TaskId], op: F) {
             for task_id in task_ids {
-                let task_id: TaskId = (*task_id).into();
-                if !op(self.get_task(task_id)) {
+                if !op(self.get_task(*task_id)) {
                     panic!("Task {task_id} does not satisfy the condition");
                 }
             }
         }
 
-        pub fn assert_worker_condition<W: Copy + Into<WorkerId>, F: Fn(&Worker) -> bool>(
+        pub fn assert_worker_condition<F: Fn(&Worker) -> bool>(
             &self,
-            worker_ids: &[W],
+            worker_ids: &[WorkerId],
             op: F,
         ) {
             for worker_id in worker_ids {
-                let worker_id: WorkerId = (*worker_id).into();
-                if !op(self.get_worker_by_id_or_panic(worker_id)) {
+                if !op(self.get_worker_by_id_or_panic(*worker_id)) {
                     panic!("Worker {worker_id} does not satisfy the condition");
                 }
             }
         }
 
-        pub fn assert_waiting<T: Copy + Into<TaskId>>(&self, task_ids: &[T]) {
+        pub fn assert_waiting(&self, task_ids: &[TaskId]) {
             self.assert_task_condition(task_ids, |t| t.is_waiting());
         }
 
-        pub fn assert_ready<T: Copy + Into<TaskId>>(&self, task_ids: &[T]) {
+        pub fn assert_ready(&self, task_ids: &[TaskId]) {
             self.assert_task_condition(task_ids, |t| t.is_ready());
         }
 
-        pub fn assert_assigned<T: Copy + Into<TaskId>>(&self, task_ids: &[T]) {
+        pub fn assert_assigned(&self, task_ids: &[TaskId]) {
             self.assert_task_condition(task_ids, |t| t.is_assigned());
         }
 
-        pub fn assert_fresh<T: Copy + Into<TaskId>>(&self, task_ids: &[T]) {
+        pub fn assert_fresh(&self, task_ids: &[TaskId]) {
             self.assert_task_condition(task_ids, |t| t.is_fresh());
         }
 
-        pub fn assert_not_fresh<T: Copy + Into<TaskId>>(&self, task_ids: &[T]) {
+        pub fn assert_not_fresh(&self, task_ids: &[TaskId]) {
             self.assert_task_condition(task_ids, |t| !t.is_fresh());
         }
 
-        pub fn assert_running<T: Copy + Into<TaskId>>(&self, task_ids: &[T]) {
+        pub fn assert_running(&self, task_ids: &[TaskId]) {
             self.assert_task_condition(task_ids, |t| t.is_sn_running());
         }
 
-        pub fn assert_underloaded<W: Copy + Into<WorkerId>>(&self, worker_ids: &[W]) {
+        pub fn assert_underloaded(&self, worker_ids: &[WorkerId]) {
             self.assert_worker_condition(worker_ids, |w| w.is_underloaded());
         }
 
-        pub fn assert_not_underloaded<W: Copy + Into<WorkerId>>(&self, worker_ids: &[W]) {
+        pub fn assert_not_underloaded(&self, worker_ids: &[WorkerId]) {
             self.assert_worker_condition(worker_ids, |w| !w.is_underloaded());
         }
 
-        pub fn task<T: Into<TaskId>>(&self, id: T) -> &Task {
-            self.tasks.get_task(id.into())
+        pub fn task(&self, id: TaskId) -> &Task {
+            self.tasks.get_task(id)
         }
     }
 

crates/tako/src/internal/server/explain.rs

@@ -160,24 +160,23 @@ mod tests {
     fn explain_single_node() {
         let mut rt = TestEnv::new();
         rt.new_named_resource("gpus");
-        rt.new_worker_with_id(1, &WorkerBuilder::new(4));
-        rt.new_worker_with_id(
-            2,
+        let w1 = rt.new_worker(&WorkerBuilder::new(4));
+        let w2 = rt.new_worker(
             &WorkerBuilder::empty()
                 .res_range("cpus", 1, 10)
                 .res_range("gpus", 1, 4)
                 .time_limit(Duration::from_secs(40_000)),
         );
 
         let resource_map = rt.core().create_resource_map();
-        let explain = |rt: &mut TestEnv, task: TaskId, worker, now| {
+        let explain = |rt: &mut TestEnv, task: TaskId, worker: WorkerId, now| {
             let group = WorkerGroup::new(Set::new());
             let (task_map, worker_map, rqs) = rt.core().split_tasks_workers_requests_mut();
             task_explain_for_worker(
                 &resource_map,
                 rqs,
                 task_map.get_task(task),
-                worker_map.get_worker(WorkerId::new(worker)),
+                worker_map.get_worker(worker),
                 &group,
                 now,
             )
@@ -187,29 +186,29 @@ mod tests {
         let now = Instant::now();
         let t1 = rt.new_task(&TaskBuilder::new());
         let (_task_map, _worker_map, _rqs) = rt.core().split_tasks_workers_requests_mut();
-        let r = explain(&mut rt, t1, 1, now);
+        let r = explain(&mut rt, t1, w1, now);
         assert_eq!(r.variants.len(), 1);
         assert_eq!(r.variants[0].len(), 1);
         assert_eq!(r.n_enabled_variants(), 1);
 
         let t2 = rt.new_task(&TaskBuilder::new().time_request(20_000));
-        let r = explain(&mut rt, t2, 1, now);
+        let r = explain(&mut rt, t2, w1, now);
         assert_eq!(r.variants.len(), 1);
         assert_eq!(r.variants[0].len(), 2);
         assert_eq!(r.n_enabled_variants(), 1);
 
-        let r = explain(&mut rt, t2, 2, now);
+        let r = explain(&mut rt, t2, w2, now);
         assert_eq!(r.variants.len(), 1);
         assert_eq!(r.variants[0].len(), 2);
         assert_eq!(r.n_enabled_variants(), 1);
 
         let now2 = now + Duration::from_secs(21_000);
-        let r = explain(&mut rt, t2, 1, now2);
+        let r = explain(&mut rt, t2, w1, now2);
         assert_eq!(r.variants.len(), 1);
         assert_eq!(r.variants[0].len(), 2);
         assert_eq!(r.n_enabled_variants(), 1);
 
-        let r = explain(&mut rt, t2, 2, now2);
+        let r = explain(&mut rt, t2, w2, now2);
         assert_eq!(r.variants.len(), 1);
         assert_eq!(r.variants[0].len(), 2);
         assert!(matches!(
@@ -227,7 +226,7 @@ mod tests {
                 .cpus(30)
                 .add_resource(1, 3),
         );
-        let r = explain(&mut rt, t3, 2, now);
+        let r = explain(&mut rt, t3, w2, now);
         assert_eq!(r.variants.len(), 1);
         assert_eq!(r.variants[0].len(), 3);
         assert!(matches!(
@@ -247,7 +246,7 @@ mod tests {
                 .cpus(2)
                 .add_resource(1, 32),
         );
-        let r = explain(&mut rt, t4, 2, now2);
+        let r = explain(&mut rt, t4, w2, now2);
         assert_eq!(r.variants.len(), 2);
         assert_eq!(r.variants[0].len(), 3);
         assert_eq!(r.variants[1].len(), 2);
@@ -283,22 +282,22 @@ mod tests {
         let now = Instant::now();
         let mut rt = TestEnv::new();
 
-        rt.new_worker_with_id(1, &WorkerBuilder::new(4));
-
+        let w1 = rt.new_worker(&WorkerBuilder::new(4));
         let t1 = rt.new_task(&TaskBuilder::new().n_nodes(4));
+
         let mut wset = Set::new();
-        wset.insert(WorkerId::new(1));
-        wset.insert(WorkerId::new(2));
-        wset.insert(WorkerId::new(3));
-        wset.insert(WorkerId::new(132));
+        wset.insert(w1);
+        wset.insert(WorkerId::new(1002));
+        wset.insert(WorkerId::new(1003));
+        wset.insert(WorkerId::new(1004));
         let group = WorkerGroup::new(wset);
         let resource_map = rt.core().create_resource_map();
         let (task_map, worker_map, rqs) = rt.core().split_tasks_workers_requests_mut();
         let r = task_explain_for_worker(
             &resource_map,
             rqs,
             task_map.get_task(t1),
-            worker_map.get_worker(WorkerId::new(1)),
+            worker_map.get_worker(w1),
             &group,
             now,
         );
@@ -307,14 +306,14 @@ mod tests {
         assert_eq!(r.n_enabled_variants(), 1);
 
         let mut wset = Set::new();
-        wset.insert(WorkerId::new(1));
-        wset.insert(WorkerId::new(132));
+        wset.insert(w1);
+        wset.insert(WorkerId::new(1032));
         let group = WorkerGroup::new(wset);
         let r = task_explain_for_worker(
             &resource_map,
             rqs,
             task_map.get_task(t1),
-            worker_map.get_worker(WorkerId::new(1)),
+            worker_map.get_worker(w1),
             &group,
             now,
         );

crates/tako/src/internal/tests/test_query.rs

@@ -1,4 +1,3 @@
-use crate::TaskId;
 use crate::control::WorkerTypeQuery;
 use crate::internal::scheduler::query::compute_new_worker_query;
 use crate::internal::server::core::Core;
@@ -92,10 +91,10 @@ fn test_query_no_enough_workers1() {
 fn test_query_enough_workers2() {
     let mut rt = TestEnv::new();
 
-    rt.new_workers_cpus(&[2]);
+    let w1 = rt.new_worker_cpus(2);
 
-    rt.new_task_running(&TaskBuilder::new(), 100);
-    rt.new_task_assigned(&TaskBuilder::new(), 100);
+    rt.new_task_running(&TaskBuilder::new(), w1);
+    rt.new_task_assigned(&TaskBuilder::new(), w1);
     rt.schedule();
 
     let r = compute_new_worker_query(
@@ -736,7 +735,7 @@ fn test_query_unknown_do_not_add_extra() {
 fn test_query_after_task_cancel() {
     let mut rt = TestEnv::new();
     let t1 = rt.new_task_cpus(10);
-    rt.new_worker_with_id(102, &WorkerBuilder::new(1));
+    rt.new_worker(&WorkerBuilder::new(1));
     rt.schedule();
     let mut comm = create_test_comm();
     on_cancel_tasks(rt.core(), &mut comm, &[t1]);