diff --git a/pvmfw/src/device_assignment.rs b/pvmfw/src/device_assignment.rs
index f826167..5edfe97 100644
--- a/pvmfw/src/device_assignment.rs
+++ b/pvmfw/src/device_assignment.rs
@@ -28,11 +28,11 @@
 use core::iter::Iterator;
 use core::mem;
 use core::ops::Range;
-// TODO(ptosi): Remove the need for this workaround.
-#[cfg(not(test))]
-use hyp::DeviceAssigningHypervisor;
 use libfdt::{Fdt, FdtError, FdtNode, FdtNodeMut, Phandle, Reg};
 use log::error;
+// TODO(b/308694211): Use vmbase::hyp::{DeviceAssigningHypervisor, Error} proper for tests.
+#[cfg(not(test))]
+use vmbase::hyp::DeviceAssigningHypervisor;
 use zerocopy::byteorder::big_endian::U32;
 use zerocopy::FromBytes as _;
 
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 72212c3..61ea65b 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -23,7 +23,6 @@
 use core::num::NonZeroUsize;
 use core::ops::Range;
 use core::slice;
-use hyp::{get_mem_sharer, get_mmio_guard};
 use log::debug;
 use log::error;
 use log::info;
@@ -32,6 +31,7 @@
 use vmbase::util::RangeExt as _;
 use vmbase::{
     configure_heap, console,
+    hyp::{get_mem_sharer, get_mmio_guard},
     layout::{self, crosvm},
     main,
     memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index 6038f12..9206588 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -46,6 +46,7 @@
 use static_assertions::const_assert;
 use tinyvec::ArrayVec;
 use vmbase::fdt::SwiotlbInfo;
+use vmbase::hyp;
 use vmbase::layout::{crosvm::MEM_START, MAX_VIRT_ADDR};
 use vmbase::memory::SIZE_4KB;
 use vmbase::util::flatten;