libavb_rs: refactor test utils

Refactor test utilities and data so that common code can be shared for
the upcoming cert tests.

Bug: b/320543206
Test: atest libavb_rs_test libavb_rs_uuid_test libavb_rs_unittest libavb_rs_uuid_unittest
Change-Id: I2fb0710bad8c2de84608c23a9092a6fc85544294

4 files changed, 208 insertions, 180 deletions

diff --git a/rust/tests/test_data.rs b/rust/tests/test_data.rs
new file mode 100644
index 0000000..f7fc499
--- /dev/null
+++ b/rust/tests/test_data.rs
@@ -0,0 +1,55 @@
+// Copyright 2024, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Test data used in libavb_rs tests.
+//!
+//! These constants must match the values used to create the images in Android.bp.
+
+pub const TEST_IMAGE_PATH: &str = "test_image.img";
+pub const TEST_IMAGE_SIZE: usize = 16 * 1024;
+pub const TEST_IMAGE_SALT_HEX: &str = "1000";
+pub const TEST_HASHTREE_SALT_HEX: &str = "B000";
+pub const TEST_VBMETA_PATH: &str = "test_vbmeta.img";
+pub const TEST_VBMETA_2_PARTITIONS_PATH: &str = "test_vbmeta_2_parts.img";
+pub const TEST_VBMETA_PERSISTENT_DIGEST_PATH: &str = "test_vbmeta_persistent_digest.img";
+pub const TEST_VBMETA_WITH_PROPERTY_PATH: &str = "test_vbmeta_with_property.img";
+pub const TEST_VBMETA_WITH_HASHTREE_PATH: &str = "test_vbmeta_with_hashtree.img";
+pub const TEST_VBMETA_WITH_COMMANDLINE_PATH: &str = "test_vbmeta_with_commandline.img";
+pub const TEST_VBMETA_WITH_CHAINED_PARTITION_PATH: &str = "test_vbmeta_with_chained_partition.img";
+pub const TEST_IMAGE_WITH_VBMETA_FOOTER_PATH: &str = "avbrs_test_image_with_vbmeta_footer.img";
+pub const TEST_IMAGE_WITH_VBMETA_FOOTER_FOR_BOOT_PATH: &str =
+    "avbrs_test_image_with_vbmeta_footer_for_boot.img";
+pub const TEST_IMAGE_WITH_VBMETA_FOOTER_FOR_TEST_PART_2: &str =
+    "avbrs_test_image_with_vbmeta_footer_for_test_part_2.img";
+pub const TEST_PUBLIC_KEY_PATH: &str = "data/testkey_rsa4096_pub.bin";
+pub const TEST_PUBLIC_KEY_RSA8192_PATH: &str = "data/testkey_rsa8192_pub.bin";
+pub const TEST_PARTITION_NAME: &str = "test_part";
+pub const TEST_PARTITION_SLOT_C_NAME: &str = "test_part_c";
+pub const TEST_PARTITION_2_NAME: &str = "test_part_2";
+pub const TEST_PARTITION_PERSISTENT_DIGEST_NAME: &str = "test_part_persistent_digest";
+pub const TEST_PARTITION_HASH_TREE_NAME: &str = "test_part_hashtree";
+pub const TEST_VBMETA_ROLLBACK_LOCATION: usize = 0; // Default value, we don't explicitly set this.
+pub const TEST_PROPERTY_KEY: &str = "test_prop_key";
+pub const TEST_PROPERTY_VALUE: &[u8] = b"test_prop_value";
+pub const TEST_KERNEL_COMMANDLINE: &str = "test_cmdline_key=test_cmdline_value";
+pub const TEST_CHAINED_PARTITION_ROLLBACK_LOCATION: usize = 4;
+pub const TEST_CHAINED_PARTITION_ROLLBACK_INDEX: u64 = 7;
+
+// Expected values determined by examining the vbmeta image with `avbtool info_image`.
+// Images can be found in <out>/soong/.intermediates/external/avb/rust/.
+pub const TEST_IMAGE_DIGEST_HEX: &str =
+    "89e6fd3142917b8c34ac7d30897a907a71bd3bf5d9b39d00bf938b41dcf3b84f";
+pub const TEST_IMAGE_HASH_ALGO: &str = "sha256";
+pub const TEST_HASHTREE_DIGEST_HEX: &str = "5373fc4ee3dd898325eeeffb5a1dbb041900c5f1";
+pub const TEST_HASHTREE_ALGORITHM: &str = "sha1";
diff --git a/rust/tests/test_ops.rs b/rust/tests/test_ops.rs
index d752650..526bf4d 100644
--- a/rust/tests/test_ops.rs
+++ b/rust/tests/test_ops.rs
@@ -66,14 +66,16 @@ impl<'a> FakePartition<'a> {
     }
 }
 
-/// Fake vbmeta key state.
-pub struct FakeVbmetaKeyState {
-    /// Key trust & rollback index info.
-    pub info: PublicKeyForPartitionInfo,
-
-    /// If specified, indicates the specific partition this vbmeta is tied to (for
-    /// `validate_public_key_for_partition()`).
-    pub for_partition: Option<&'static str>,
+/// Fake vbmeta key.
+pub enum FakeVbmetaKey {
+    /// Standard AVB validation using a hardcoded key; if the signing key matches these contents
+    /// it is accepted, otherwise it's rejected.
+    Avb {
+        /// Expected public key contents.
+        public_key: Vec<u8>,
+        /// Expected public key metadata contents.
+        public_key_metadata: Option<Vec<u8>>,
+    },
 }
 
 /// Fake `Ops` test fixture.
@@ -84,11 +86,15 @@ pub struct TestOps<'a> {
     /// Partitions to provide to libavb callbacks.
     pub partitions: HashMap<&'static str, FakePartition<'a>>,
 
-    /// Vbmeta public keys as a map of {(key, metadata): state}. Querying unknown keys will
-    /// return `IoError::Io`.
+    /// Default vbmeta key to use for the `validate_vbmeta_public_key()` callback, or `None` to
+    /// return `IoError::Io` when accessing this key.
+    pub default_vbmeta_key: Option<FakeVbmetaKey>,
+
+    /// Additional vbmeta keys for the `validate_public_key_for_partition()` callback.
     ///
-    /// See `add_vbmeta_key*()` functions for simpler wrappers to inject these keys.
-    pub vbmeta_keys: HashMap<(Vec<u8>, Option<Vec<u8>>), FakeVbmetaKeyState>,
+    /// Stored as a map of {partition_name: (key, rollback_location)}. Querying keys for partitions
+    /// not in this map will return `IoError::Io`.
+    pub vbmeta_keys_for_partition: HashMap<&'static str, (FakeVbmetaKey, u32)>,
 
     /// Rollback indices. Accessing unknown locations will return `IoError::Io`.
     pub rollbacks: HashMap<usize, u64>,
@@ -153,41 +159,29 @@ impl<'a> TestOps<'a> {
             .insert(name.into(), contents.map(|b| b.into()));
     }
 
-    /// Adds a fake vbmeta key not tied to any partition.
-    pub fn add_vbmeta_key(&mut self, key: Vec<u8>, metadata: Option<Vec<u8>>, trusted: bool) {
-        self.vbmeta_keys.insert(
-            (key, metadata),
-            FakeVbmetaKeyState {
-                // `rollback_index_location` doesn't matter in this case, it will be read from
-                // the vbmeta blob.
-                info: PublicKeyForPartitionInfo {
-                    trusted,
-                    rollback_index_location: 0,
-                },
-                for_partition: None,
-            },
-        );
-    }
-
-    /// Adds a fake vbmeta key tied to the given partition and rollback index location.
-    pub fn add_vbmeta_key_for_partition(
+    /// Internal helper to validate a vbmeta key.
+    fn validate_fake_key(
         &mut self,
-        key: Vec<u8>,
-        metadata: Option<Vec<u8>>,
-        trusted: bool,
-        partition: &'static str,
-        rollback_index_location: u32,
-    ) {
-        self.vbmeta_keys.insert(
-            (key, metadata),
-            FakeVbmetaKeyState {
-                info: PublicKeyForPartitionInfo {
-                    trusted,
-                    rollback_index_location,
-                },
-                for_partition: Some(partition),
-            },
-        );
+        partition: Option<&str>,
+        public_key: &[u8],
+        public_key_metadata: Option<&[u8]>,
+    ) -> IoResult<bool> {
+        let fake_key = match partition {
+            None => self.default_vbmeta_key.as_ref(),
+            Some(p) => self.vbmeta_keys_for_partition.get(p).map(|(key, _)| key),
+        }
+        .ok_or(IoError::Io)?;
+
+        match fake_key {
+            FakeVbmetaKey::Avb {
+                public_key: expected_key,
+                public_key_metadata: expected_metadata,
+            } => {
+                // avb: only accept if it matches the hardcoded key + metadata.
+                Ok(expected_key == public_key
+                    && expected_metadata.as_deref() == public_key_metadata)
+            }
+        }
     }
 }
 
@@ -195,7 +189,8 @@ impl Default for TestOps<'_> {
     fn default() -> Self {
         Self {
             partitions: HashMap::new(),
-            vbmeta_keys: HashMap::new(),
+            default_vbmeta_key: None,
+            vbmeta_keys_for_partition: HashMap::new(),
             rollbacks: HashMap::new(),
             unlock_state: Err(IoError::Io),
             persistent_values: HashMap::new(),
@@ -259,13 +254,7 @@ impl<'a> Ops<'a> for TestOps<'a> {
         public_key: &[u8],
         public_key_metadata: Option<&[u8]>,
     ) -> IoResult<bool> {
-        self.vbmeta_keys
-            // The compiler can't match (&[u8], Option<&[u8]>) to keys of type
-            // (Vec<u8>, Option<Vec<u8>>) so we turn the &[u8] into vectors here. This is a bit
-            // inefficient, but it's simple which is more important for tests than efficiency.
-            .get(&(public_key.to_vec(), public_key_metadata.map(|m| m.to_vec())))
-            .ok_or(IoError::Io)
-            .map(|k| k.info.trusted)
+        self.validate_fake_key(None, public_key, public_key_metadata)
     }
 
     fn read_rollback_index(&mut self, location: usize) -> IoResult<u64> {
@@ -345,19 +334,17 @@ impl<'a> Ops<'a> for TestOps<'a> {
         public_key: &[u8],
         public_key_metadata: Option<&[u8]>,
     ) -> IoResult<PublicKeyForPartitionInfo> {
-        let key = self
-            .vbmeta_keys
-            .get(&(public_key.to_vec(), public_key_metadata.map(|m| m.to_vec())))
-            .ok_or(IoError::Io)?;
-
-        if let Some(for_partition) = key.for_partition {
-            if for_partition == partition.to_str()? {
-                // The key is registered for this partition; return its info.
-                return Ok(key.info);
-            }
-        }
-
-        // No match.
-        Err(IoError::Io)
+        let partition = partition.to_str()?;
+
+        let rollback_index_location = self
+            .vbmeta_keys_for_partition
+            .get(partition)
+            .ok_or(IoError::Io)?
+            .1;
+
+        Ok(PublicKeyForPartitionInfo {
+            trusted: self.validate_fake_key(Some(partition), public_key, public_key_metadata)?,
+            rollback_index_location,
+        })
     }
 }
diff --git a/rust/tests/tests.rs b/rust/tests/tests.rs
index 0e9b3e2..157ea4f 100644
--- a/rust/tests/tests.rs
+++ b/rust/tests/tests.rs
@@ -12,5 +12,41 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod test_data;
 mod test_ops;
 mod verify_tests;
+
+use avb::{slot_verify, HashtreeErrorMode, SlotVerifyData, SlotVerifyFlags, SlotVerifyResult};
+use std::{ffi::CString, fs};
+use test_data::*;
+use test_ops::{FakeVbmetaKey, TestOps};
+
+/// Initializes a `TestOps` object such that verification will succeed on `TEST_PARTITION_NAME`.
+///
+/// This usually forms the basis of the `TestOps` objects used, with tests modifying the returned
+/// object as needed for the individual test case.
+fn build_test_ops_one_image_one_vbmeta<'a>() -> TestOps<'a> {
+    let mut ops = TestOps::default();
+    ops.add_partition(TEST_PARTITION_NAME, fs::read(TEST_IMAGE_PATH).unwrap());
+    ops.add_partition("vbmeta", fs::read(TEST_VBMETA_PATH).unwrap());
+    ops.default_vbmeta_key = Some(FakeVbmetaKey::Avb {
+        public_key: fs::read(TEST_PUBLIC_KEY_PATH).unwrap(),
+        public_key_metadata: None,
+    });
+    ops.rollbacks.insert(TEST_VBMETA_ROLLBACK_LOCATION, 0);
+    ops.unlock_state = Ok(false);
+    ops
+}
+
+/// Calls `slot_verify()` using standard args for `build_test_ops_one_image_one_vbmeta()` setup.
+fn verify_one_image_one_vbmeta<'a>(
+    ops: &mut TestOps<'a>,
+) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
+    slot_verify(
+        ops,
+        &[&CString::new(TEST_PARTITION_NAME).unwrap()],
+        None,
+        SlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NONE,
+        HashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_EIO,
+    )
+}
diff --git a/rust/tests/verify_tests.rs b/rust/tests/verify_tests.rs
index 2e4aaad..dfe4ec8 100644
--- a/rust/tests/verify_tests.rs
+++ b/rust/tests/verify_tests.rs
@@ -14,7 +14,12 @@
 
 //! libavb_rs verification tests.
 
-use crate::test_ops::TestOps;
+use crate::{
+    build_test_ops_one_image_one_vbmeta,
+    test_data::*,
+    test_ops::{FakeVbmetaKey, TestOps},
+    verify_one_image_one_vbmeta,
+};
 use avb::{
     slot_verify, ChainPartitionDescriptor, ChainPartitionDescriptorFlags, Descriptor,
     HashDescriptor, HashDescriptorFlags, HashtreeDescriptor, HashtreeDescriptorFlags,
@@ -26,73 +31,10 @@ use std::{ffi::CString, fs};
 #[cfg(feature = "uuid")]
 use uuid::uuid;
 
-// These constants must match the values used to create the images in Android.bp.
-const TEST_IMAGE_PATH: &str = "test_image.img";
-const TEST_IMAGE_SIZE: usize = 16 * 1024;
-const TEST_IMAGE_SALT_HEX: &str = "1000";
-const TEST_HASHTREE_SALT_HEX: &str = "B000";
-const TEST_VBMETA_PATH: &str = "test_vbmeta.img";
-const TEST_VBMETA_2_PARTITIONS_PATH: &str = "test_vbmeta_2_parts.img";
-const TEST_VBMETA_PERSISTENT_DIGEST_PATH: &str = "test_vbmeta_persistent_digest.img";
-const TEST_VBMETA_WITH_PROPERTY_PATH: &str = "test_vbmeta_with_property.img";
-const TEST_VBMETA_WITH_HASHTREE_PATH: &str = "test_vbmeta_with_hashtree.img";
-const TEST_VBMETA_WITH_COMMANDLINE_PATH: &str = "test_vbmeta_with_commandline.img";
-const TEST_VBMETA_WITH_CHAINED_PARTITION_PATH: &str = "test_vbmeta_with_chained_partition.img";
-const TEST_IMAGE_WITH_VBMETA_FOOTER_PATH: &str = "avbrs_test_image_with_vbmeta_footer.img";
-const TEST_IMAGE_WITH_VBMETA_FOOTER_FOR_BOOT_PATH: &str =
-    "avbrs_test_image_with_vbmeta_footer_for_boot.img";
-const TEST_IMAGE_WITH_VBMETA_FOOTER_FOR_TEST_PART_2: &str =
-    "avbrs_test_image_with_vbmeta_footer_for_test_part_2.img";
-const TEST_PUBLIC_KEY_PATH: &str = "data/testkey_rsa4096_pub.bin";
-const TEST_PUBLIC_KEY_RSA8192_PATH: &str = "data/testkey_rsa8192_pub.bin";
-const TEST_PARTITION_NAME: &str = "test_part";
-const TEST_PARTITION_SLOT_C_NAME: &str = "test_part_c";
-const TEST_PARTITION_2_NAME: &str = "test_part_2";
-const TEST_PARTITION_PERSISTENT_DIGEST_NAME: &str = "test_part_persistent_digest";
-const TEST_PARTITION_HASH_TREE_NAME: &str = "test_part_hashtree";
-const TEST_VBMETA_ROLLBACK_LOCATION: usize = 0; // Default value, we don't explicitly set this.
-const TEST_PROPERTY_KEY: &str = "test_prop_key";
-const TEST_PROPERTY_VALUE: &[u8] = b"test_prop_value";
-const TEST_KERNEL_COMMANDLINE: &str = "test_cmdline_key=test_cmdline_value";
-const TEST_CHAINED_PARTITION_ROLLBACK_LOCATION: usize = 4;
-const TEST_CHAINED_PARTITION_ROLLBACK_INDEX: u64 = 7;
-
-// Expected values determined by examining the vbmeta image with `avbtool info_image`.
-// Images can be found in <out>/soong/.intermediates/external/avb/rust/.
-const TEST_IMAGE_DIGEST_HEX: &str =
-    "89e6fd3142917b8c34ac7d30897a907a71bd3bf5d9b39d00bf938b41dcf3b84f";
-const TEST_IMAGE_HASH_ALGO: &str = "sha256";
-const TEST_HASHTREE_DIGEST_HEX: &str = "5373fc4ee3dd898325eeeffb5a1dbb041900c5f1";
-const TEST_HASHTREE_ALGORITHM: &str = "sha1";
-
-/// Initializes a `TestOps` object such that verification will succeed on `TEST_PARTITION_NAME`.
-fn test_ops_one_image_one_vbmeta<'a>() -> TestOps<'a> {
-    let mut ops = TestOps::default();
-    ops.add_partition(TEST_PARTITION_NAME, fs::read(TEST_IMAGE_PATH).unwrap());
-    ops.add_partition("vbmeta", fs::read(TEST_VBMETA_PATH).unwrap());
-    ops.add_vbmeta_key(fs::read(TEST_PUBLIC_KEY_PATH).unwrap(), None, true);
-    ops.rollbacks.insert(TEST_VBMETA_ROLLBACK_LOCATION, 0);
-    ops.unlock_state = Ok(false);
-    ops
-}
-
-/// Calls `slot_verify()` using standard args for `test_ops_one_image_one_vbmeta()` setup.
-fn verify_one_image_one_vbmeta<'a>(
-    ops: &mut TestOps<'a>,
-) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
-    slot_verify(
-        ops,
-        &[&CString::new(TEST_PARTITION_NAME).unwrap()],
-        None,
-        SlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NONE,
-        HashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_EIO,
-    )
-}
-
 /// Initializes a `TestOps` object such that verification will succeed on `TEST_PARTITION_NAME` and
 /// `TEST_PARTITION_2_NAME`.
-fn test_ops_two_images_one_vbmeta<'a>() -> TestOps<'a> {
-    let mut ops = test_ops_one_image_one_vbmeta();
+fn build_test_ops_two_images_one_vbmeta<'a>() -> TestOps<'a> {
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     // Add in the contents of the second partition and overwrite the vbmeta partition to
     // include both partition descriptors.
     ops.add_partition(TEST_PARTITION_2_NAME, fs::read(TEST_IMAGE_PATH).unwrap());
@@ -116,8 +58,8 @@ fn verify_two_images<'a>(ops: &mut TestOps<'a>) -> SlotVerifyResult<'a, SlotVeri
 
 /// Initializes a `TestOps` object such that verification will succeed on the `boot` partition with
 /// a combined image + vbmeta.
-fn test_ops_boot_partition<'a>() -> TestOps<'a> {
-    let mut ops = test_ops_one_image_one_vbmeta();
+fn build_test_ops_boot_partition<'a>() -> TestOps<'a> {
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.partitions.clear();
     ops.add_partition(
         "boot",
@@ -126,7 +68,7 @@ fn test_ops_boot_partition<'a>() -> TestOps<'a> {
     ops
 }
 
-/// Calls `slot_verify()` using standard args for `test_ops_boot_partition()` setup.
+/// Calls `slot_verify()` using standard args for `build_test_ops_boot_partition()` setup.
 fn verify_boot_partition<'a>(ops: &mut TestOps<'a>) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
     slot_verify(
         ops,
@@ -142,8 +84,8 @@ fn verify_boot_partition<'a>(ops: &mut TestOps<'a>) -> SlotVerifyResult<'a, Slot
 
 /// Initializes a `TestOps` object such that verification will succeed on
 /// `TEST_PARTITION_PERSISTENT_DIGEST_NAME`.
-fn test_ops_persistent_digest<'a>(image: Vec<u8>) -> TestOps<'a> {
-    let mut ops = test_ops_one_image_one_vbmeta();
+fn build_test_ops_persistent_digest<'a>(image: Vec<u8>) -> TestOps<'a> {
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.partitions.clear();
     // Use the vbmeta image with the persistent digest descriptor.
     ops.add_partition(
@@ -155,7 +97,7 @@ fn test_ops_persistent_digest<'a>(image: Vec<u8>) -> TestOps<'a> {
     ops
 }
 
-/// Calls `slot_verify()` using standard args for `test_ops_persistent_digest()` setup.
+/// Calls `slot_verify()` using standard args for `build_test_ops_persistent_digest()` setup.
 fn verify_persistent_digest<'a>(ops: &mut TestOps<'a>) -> SlotVerifyResult<'a, SlotVerifyData<'a>> {
     slot_verify(
         ops,
@@ -184,7 +126,7 @@ fn persistent_digest_value_name() -> String {
 
 #[test]
 fn one_image_one_vbmeta_passes_verification_with_correct_data() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
 
     let result = verify_one_image_one_vbmeta(&mut ops);
 
@@ -225,7 +167,7 @@ fn one_image_one_vbmeta_passes_verification_with_correct_data() {
 
 #[test]
 fn preloaded_image_passes_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     // Use preloaded data instead for the test partition.
     let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
     ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
@@ -243,7 +185,7 @@ fn preloaded_image_passes_verification() {
 #[test]
 fn verification_data_from_disk_can_outlive_ops() {
     let result = {
-        let mut ops = test_ops_one_image_one_vbmeta();
+        let mut ops = build_test_ops_one_image_one_vbmeta();
         verify_one_image_one_vbmeta(&mut ops)
     };
 
@@ -264,7 +206,7 @@ fn verification_data_preloaded_can_outlive_ops() {
     let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
 
     let result = {
-        let mut ops = test_ops_one_image_one_vbmeta();
+        let mut ops = build_test_ops_one_image_one_vbmeta();
         ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
         verify_one_image_one_vbmeta(&mut ops)
     };
@@ -287,7 +229,7 @@ fn verification_data_preloaded_can_outlive_ops() {
 // fn verification_data_preloaded_cannot_outlive_result() {
 //     let result = {
 //         let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
-//         let mut ops = test_ops_one_image_one_vbmeta();
+//         let mut ops = build_test_ops_one_image_one_vbmeta();
 //         ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
 //         verify_one_image_one_vbmeta(&mut ops)
 //     };
@@ -296,7 +238,7 @@ fn verification_data_preloaded_can_outlive_ops() {
 
 #[test]
 fn slotted_partition_passes_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     // Move the partitions to a "_c" slot.
     ops.partitions.clear();
     ops.add_partition(
@@ -319,7 +261,7 @@ fn slotted_partition_passes_verification() {
 
 #[test]
 fn two_images_one_vbmeta_passes_verification() {
-    let mut ops = test_ops_two_images_one_vbmeta();
+    let mut ops = build_test_ops_two_images_one_vbmeta();
 
     let result = verify_two_images(&mut ops);
 
@@ -340,20 +282,25 @@ fn two_images_one_vbmeta_passes_verification() {
 
 #[test]
 fn combined_image_vbmeta_partition_passes_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.partitions.clear();
     // Register the single combined image + vbmeta in `TEST_PARTITION_NAME`.
     ops.add_partition(
         TEST_PARTITION_NAME,
         fs::read(TEST_IMAGE_WITH_VBMETA_FOOTER_PATH).unwrap(),
     );
-    // For a combined image we need to register the public key specifically for this partition.
-    ops.add_vbmeta_key_for_partition(
-        fs::read(TEST_PUBLIC_KEY_PATH).unwrap(),
-        None,
-        true,
+    // For a combined image it should not attempt to use the default "vbmeta" key, instead we
+    // register the public key specifically for this partition.
+    ops.default_vbmeta_key = None;
+    ops.vbmeta_keys_for_partition.insert(
         TEST_PARTITION_NAME,
-        TEST_VBMETA_ROLLBACK_LOCATION as u32,
+        (
+            FakeVbmetaKey::Avb {
+                public_key: fs::read(TEST_PUBLIC_KEY_PATH).unwrap(),
+                public_key_metadata: None,
+            },
+            TEST_VBMETA_ROLLBACK_LOCATION as u32,
+        ),
     );
 
     let result = slot_verify(
@@ -389,7 +336,7 @@ fn combined_image_vbmeta_partition_passes_verification() {
 // Validate the custom behavior if the combined image + vbmeta live in the `boot` partition.
 #[test]
 fn vbmeta_with_boot_partition_passes_verification() {
-    let mut ops = test_ops_boot_partition();
+    let mut ops = build_test_ops_boot_partition();
 
     let result = verify_boot_partition(&mut ops);
 
@@ -413,7 +360,7 @@ fn persistent_digest_verification_updates_persistent_value() {
     // calculated on-demand and stored into a named persistent value. So our test image can contain
     // anything, but does have to match the size indicated by the descriptor.
     let image_contents = vec![0xAAu8; TEST_IMAGE_SIZE];
-    let mut ops = test_ops_persistent_digest(image_contents.clone());
+    let mut ops = build_test_ops_persistent_digest(image_contents.clone());
 
     {
         let result = verify_persistent_digest(&mut ops);
@@ -429,7 +376,7 @@ fn persistent_digest_verification_updates_persistent_value() {
 #[cfg(feature = "uuid")]
 #[test]
 fn successful_verification_substitutes_partition_guid() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.partitions.get_mut("vbmeta").unwrap().uuid = uuid!("01234567-89ab-cdef-0123-456789abcdef");
 
     let result = verify_one_image_one_vbmeta(&mut ops);
@@ -445,7 +392,7 @@ fn successful_verification_substitutes_partition_guid() {
 #[cfg(feature = "uuid")]
 #[test]
 fn successful_verification_substitutes_boot_partition_guid() {
-    let mut ops = test_ops_boot_partition();
+    let mut ops = build_test_ops_boot_partition();
     ops.partitions.get_mut("boot").unwrap().uuid = uuid!("01234567-89ab-cdef-0123-456789abcdef");
 
     let result = verify_boot_partition(&mut ops);
@@ -461,7 +408,7 @@ fn successful_verification_substitutes_boot_partition_guid() {
 
 #[test]
 fn corrupted_image_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     modify_partition_contents(&mut ops, TEST_PARTITION_NAME);
 
     let result = verify_one_image_one_vbmeta(&mut ops);
@@ -472,7 +419,7 @@ fn corrupted_image_fails_verification() {
 
 #[test]
 fn read_partition_callback_error_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.partitions.remove(TEST_PARTITION_NAME);
 
     let result = verify_one_image_one_vbmeta(&mut ops);
@@ -483,7 +430,7 @@ fn read_partition_callback_error_fails_verification() {
 
 #[test]
 fn undersized_partition_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.partitions
         .get_mut(TEST_PARTITION_NAME)
         .unwrap()
@@ -499,7 +446,7 @@ fn undersized_partition_fails_verification() {
 
 #[test]
 fn corrupted_vbmeta_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     modify_partition_contents(&mut ops, "vbmeta");
 
     let result = verify_one_image_one_vbmeta(&mut ops);
@@ -510,7 +457,7 @@ fn corrupted_vbmeta_fails_verification() {
 
 #[test]
 fn rollback_violation_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     // Device with rollback = 1 should refuse to boot image with rollback = 0.
     ops.rollbacks.insert(TEST_VBMETA_ROLLBACK_LOCATION, 1);
 
@@ -522,7 +469,7 @@ fn rollback_violation_fails_verification() {
 
 #[test]
 fn rollback_callback_error_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.rollbacks.clear();
 
     let result = verify_one_image_one_vbmeta(&mut ops);
@@ -533,8 +480,11 @@ fn rollback_callback_error_fails_verification() {
 
 #[test]
 fn untrusted_vbmeta_keys_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
-    ops.add_vbmeta_key(fs::read(TEST_PUBLIC_KEY_PATH).unwrap(), None, false);
+    let mut ops = build_test_ops_one_image_one_vbmeta();
+    ops.default_vbmeta_key = Some(FakeVbmetaKey::Avb {
+        public_key: b"not_the_key".into(),
+        public_key_metadata: None,
+    });
 
     let result = verify_one_image_one_vbmeta(&mut ops);
 
@@ -544,8 +494,8 @@ fn untrusted_vbmeta_keys_fails_verification() {
 
 #[test]
 fn vbmeta_keys_callback_error_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
-    ops.vbmeta_keys.clear();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
+    ops.default_vbmeta_key = None;
 
     let result = verify_one_image_one_vbmeta(&mut ops);
 
@@ -555,7 +505,7 @@ fn vbmeta_keys_callback_error_fails_verification() {
 
 #[test]
 fn unlock_state_callback_error_fails_verification() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     ops.unlock_state = Err(IoError::Io);
 
     let result = verify_one_image_one_vbmeta(&mut ops);
@@ -567,7 +517,7 @@ fn unlock_state_callback_error_fails_verification() {
 #[test]
 fn persistent_digest_mismatch_fails_verification() {
     let image_contents = vec![0xAAu8; TEST_IMAGE_SIZE];
-    let mut ops = test_ops_persistent_digest(image_contents.clone());
+    let mut ops = build_test_ops_persistent_digest(image_contents.clone());
     // Put in an incorrect persistent digest; `slot_verify()` should detect the mismatch and fail.
     ops.add_persistent_value(&persistent_digest_value_name(), Ok(b"incorrect_digest"));
     // Make a copy so we can verify the persistent values don't change on failure.
@@ -582,7 +532,7 @@ fn persistent_digest_mismatch_fails_verification() {
 #[test]
 fn persistent_digest_callback_error_fails_verification() {
     let image_contents = vec![0xAAu8; TEST_IMAGE_SIZE];
-    let mut ops = test_ops_persistent_digest(image_contents.clone());
+    let mut ops = build_test_ops_persistent_digest(image_contents.clone());
     ops.add_persistent_value(&persistent_digest_value_name(), Err(IoError::NoSuchValue));
 
     let result = verify_persistent_digest(&mut ops);
@@ -593,7 +543,7 @@ fn persistent_digest_callback_error_fails_verification() {
 
 #[test]
 fn corrupted_image_with_allow_verification_error_flag_fails_verification_with_data() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     modify_partition_contents(&mut ops, TEST_PARTITION_NAME);
 
     let result = slot_verify(
@@ -626,7 +576,7 @@ fn corrupted_image_with_allow_verification_error_flag_fails_verification_with_da
 
 #[test]
 fn one_image_one_vbmeta_verification_data_display() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
 
     let result = verify_one_image_one_vbmeta(&mut ops);
 
@@ -639,7 +589,7 @@ fn one_image_one_vbmeta_verification_data_display() {
 
 #[test]
 fn preloaded_image_verification_data_display() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     let preloaded = fs::read(TEST_IMAGE_PATH).unwrap();
     ops.add_preloaded_partition(TEST_PARTITION_NAME, &preloaded);
 
@@ -654,7 +604,7 @@ fn preloaded_image_verification_data_display() {
 
 #[test]
 fn two_images_one_vbmeta_verification_data_display() {
-    let mut ops = test_ops_two_images_one_vbmeta();
+    let mut ops = build_test_ops_two_images_one_vbmeta();
 
     let result = verify_two_images(&mut ops);
 
@@ -667,7 +617,7 @@ fn two_images_one_vbmeta_verification_data_display() {
 
 #[test]
 fn corrupted_image_verification_data_display() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
     modify_partition_contents(&mut ops, TEST_PARTITION_NAME);
 
     let result = slot_verify(
@@ -691,7 +641,7 @@ fn corrupted_image_verification_data_display() {
 
 #[test]
 fn one_image_gives_single_descriptor() {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
 
     let result = verify_one_image_one_vbmeta(&mut ops);
 
@@ -701,7 +651,7 @@ fn one_image_gives_single_descriptor() {
 
 #[test]
 fn two_images_gives_two_descriptors() {
-    let mut ops = test_ops_two_images_one_vbmeta();
+    let mut ops = build_test_ops_two_images_one_vbmeta();
 
     let result = verify_two_images(&mut ops);
 
@@ -718,7 +668,7 @@ fn two_images_gives_two_descriptors() {
 /// 3. run verification
 /// 4. check that the given `descriptor` exists in the verification data
 fn verify_and_find_descriptor(vbmeta_path: &str, expected_descriptor: &Descriptor) {
-    let mut ops = test_ops_one_image_one_vbmeta();
+    let mut ops = build_test_ops_one_image_one_vbmeta();
 
     // Replace the vbmeta image with the requested variation.
     ops.add_partition("vbmeta", fs::read(vbmeta_path).unwrap());
@@ -793,7 +743,7 @@ fn verify_kernel_commandline_descriptor() {
 
 #[test]
 fn verify_chain_partition_descriptor() {
-    let mut ops = test_ops_two_images_one_vbmeta();
+    let mut ops = build_test_ops_two_images_one_vbmeta();
 
     // Set up the fake ops to contain:
     // * the default test image in TEST_PARTITION_NAME