1 files changed, 15 insertions, 10 deletions
diff --git a/nitrocli/src/tests/unencrypted.rs b/nitrocli/src/tests/unencrypted.rs
index c976f50..547dcaf 100644
--- a/nitrocli/src/tests/unencrypted.rs
+++ b/nitrocli/src/tests/unencrypted.rs
@@ -19,23 +19,28 @@
 
 use super::*;
 
-#[test_device]
-fn unencrypted_set_read_write(device: nitrokey::Storage) -> crate::Result<()> {
-  let mut ncli = Nitrocli::with_dev(device);
+#[test_device(storage)]
+fn unencrypted_set_read_write(model: nitrokey::Model) -> crate::Result<()> {
+  let mut ncli = Nitrocli::with_model(model);
   let out = ncli.handle(&["unencrypted", "set", "read-write"])?;
   assert!(out.is_empty());
 
-  let device = nitrokey::Storage::connect()?;
-  assert!(device.get_status()?.unencrypted_volume.active);
-  assert!(!device.get_status()?.unencrypted_volume.read_only);
-  drop(device);
+  {
+    let mut manager = nitrokey::force_take()?;
+    let device = manager.connect_storage()?;
+    assert!(device.get_status()?.unencrypted_volume.active);
+    assert!(!device.get_status()?.unencrypted_volume.read_only);
+  }
 
   let out = ncli.handle(&["unencrypted", "set", "read-only"])?;
   assert!(out.is_empty());
 
-  let device = nitrokey::Storage::connect()?;
-  assert!(device.get_status()?.unencrypted_volume.active);
-  assert!(device.get_status()?.unencrypted_volume.read_only);
+  {
+    let mut manager = nitrokey::force_take()?;
+    let device = manager.connect_storage()?;
+    assert!(device.get_status()?.unencrypted_volume.active);
+    assert!(device.get_status()?.unencrypted_volume.read_only);
+  }
 
   Ok(())
 }