fix helios tests

oxidecomputer · Jul 1, 2024 · 2c6eb01 · 2c6eb01
1 parent d6618e7
commit 2c6eb01
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/sled-storage/src/manager.rs b/sled-storage/src/manager.rs
@@ -1083,8 +1083,16 @@ mod tests {
 
         // Now let's verify we saw the correct firmware update.
         for rd in &raw_disks {
-            let managed =
-                all_disks_gen2.values.get(rd.identity()).expect("disk exists");
+            let managed = all_disks_gen2
+                .iter_all_inner()
+                .find_map(|(identity, disk)| {
+                    if identity == rd.identity() {
+                        Some(disk)
+                    } else {
+                        None
+                    }
+                })
+                .expect("disk exists");
             match managed {
                 ManagedDisk::ExplicitlyManaged(disk)
                 | ManagedDisk::ImplicitlyManaged(disk) => {
@@ -1321,7 +1329,7 @@ mod tests {
         let expected: HashSet<_> =
             disks.iter().skip(1).take(3).map(|d| d.identity()).collect();
         let actual: HashSet<_> =
-            all_disks.iter_all().map(|(identity, _, _)| identity).collect();
+            all_disks.iter_all().map(|(identity, _, _, _)| identity).collect();
         assert_eq!(expected, actual);
 
         // Ensure the same set of disks and make sure no change occurs
@@ -1355,7 +1363,7 @@ mod tests {
         let expected: HashSet<_> =
             disks.iter().skip(4).take(5).map(|d| d.identity()).collect();
         let actual: HashSet<_> =
-            all_disks.iter_all().map(|(identity, _, _)| identity).collect();
+            all_disks.iter_all().map(|(identity, _, _, _)| identity).collect();
         assert_eq!(expected, actual);
 
         harness.cleanup().await;