Move large test modules into separate files

This complies better with the Android Rust style guide, and makes it
easier to navigate the code.

Test: keystore2_test  libwatchdog_rs.test  librkpd_client.test
Change-Id: Iceb49e309af66ec16d31da66b328936b0312061a
diff --git a/keystore2/rkpd_client/src/lib.rs b/keystore2/rkpd_client/src/lib.rs
index d8a5276..936fe3d 100644
--- a/keystore2/rkpd_client/src/lib.rs
+++ b/keystore2/rkpd_client/src/lib.rs
@@ -310,331 +310,4 @@
 }
 
 #[cfg(test)]
-mod tests {
-    use super::*;
-    use android_security_rkp_aidl::aidl::android::security::rkp::IRegistration::BnRegistration;
-    use std::sync::atomic::{AtomicU32, Ordering};
-    use std::sync::{Arc, Mutex};
-
-    const DEFAULT_RPC_SERVICE_NAME: &str =
-        "android.hardware.security.keymint.IRemotelyProvisionedComponent/default";
-
-    struct MockRegistrationValues {
-        key: RemotelyProvisionedKey,
-        latency: Option<Duration>,
-        thread_join_handles: Vec<Option<std::thread::JoinHandle<()>>>,
-    }
-
-    struct MockRegistration(Arc<Mutex<MockRegistrationValues>>);
-
-    impl MockRegistration {
-        pub fn new_native_binder(
-            key: &RemotelyProvisionedKey,
-            latency: Option<Duration>,
-        ) -> Strong<dyn IRegistration> {
-            let result = Self(Arc::new(Mutex::new(MockRegistrationValues {
-                key: RemotelyProvisionedKey {
-                    keyBlob: key.keyBlob.clone(),
-                    encodedCertChain: key.encodedCertChain.clone(),
-                },
-                latency,
-                thread_join_handles: Vec::new(),
-            })));
-            BnRegistration::new_binder(result, BinderFeatures::default())
-        }
-    }
-
-    impl Drop for MockRegistration {
-        fn drop(&mut self) {
-            let mut values = self.0.lock().unwrap();
-            for handle in values.thread_join_handles.iter_mut() {
-                // These are test threads. So, no need to worry too much about error handling.
-                handle.take().unwrap().join().unwrap();
-            }
-        }
-    }
-
-    impl Interface for MockRegistration {}
-
-    impl IRegistration for MockRegistration {
-        fn getKey(&self, _: i32, cb: &Strong<dyn IGetKeyCallback>) -> binder::Result<()> {
-            let mut values = self.0.lock().unwrap();
-            let key = RemotelyProvisionedKey {
-                keyBlob: values.key.keyBlob.clone(),
-                encodedCertChain: values.key.encodedCertChain.clone(),
-            };
-            let latency = values.latency;
-            let get_key_cb = cb.clone();
-
-            // Need a separate thread to trigger timeout in the caller.
-            let join_handle = std::thread::spawn(move || {
-                if let Some(duration) = latency {
-                    std::thread::sleep(duration);
-                }
-                get_key_cb.onSuccess(&key).unwrap();
-            });
-            values.thread_join_handles.push(Some(join_handle));
-            Ok(())
-        }
-
-        fn cancelGetKey(&self, _: &Strong<dyn IGetKeyCallback>) -> binder::Result<()> {
-            Ok(())
-        }
-
-        fn storeUpgradedKeyAsync(
-            &self,
-            _: &[u8],
-            _: &[u8],
-            cb: &Strong<dyn IStoreUpgradedKeyCallback>,
-        ) -> binder::Result<()> {
-            // We are primarily concerned with timing out correctly. Storing the key in this mock
-            // registration isn't particularly interesting, so skip that part.
-            let values = self.0.lock().unwrap();
-            let store_cb = cb.clone();
-            let latency = values.latency;
-
-            std::thread::spawn(move || {
-                if let Some(duration) = latency {
-                    std::thread::sleep(duration);
-                }
-                store_cb.onSuccess().unwrap();
-            });
-            Ok(())
-        }
-    }
-
-    fn get_mock_registration(
-        key: &RemotelyProvisionedKey,
-        latency: Option<Duration>,
-    ) -> Result<binder::Strong<dyn IRegistration>> {
-        let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx);
-        let mock_registration = MockRegistration::new_native_binder(key, latency);
-
-        assert!(cb.onSuccess(&mock_registration).is_ok());
-        tokio_rt().block_on(rx).unwrap()
-    }
-
-    // Using the same key ID makes test cases race with each other. So, we use separate key IDs for
-    // different test cases.
-    fn get_next_key_id() -> u32 {
-        static ID: AtomicU32 = AtomicU32::new(0);
-        ID.fetch_add(1, Ordering::Relaxed)
-    }
-
-    #[test]
-    fn test_get_registration_cb_success() {
-        let key: RemotelyProvisionedKey = Default::default();
-        let registration = get_mock_registration(&key, /*latency=*/ None);
-        assert!(registration.is_ok());
-    }
-
-    #[test]
-    fn test_get_registration_cb_cancel() {
-        let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx);
-        assert!(cb.onCancel().is_ok());
-
-        let result = tokio_rt().block_on(rx).unwrap();
-        assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::RequestCancelled);
-    }
-
-    #[test]
-    fn test_get_registration_cb_error() {
-        let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx);
-        assert!(cb.onError("error").is_ok());
-
-        let result = tokio_rt().block_on(rx).unwrap();
-        assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::GetRegistrationFailed);
-    }
-
-    #[test]
-    fn test_get_key_cb_success() {
-        let mock_key =
-            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
-        let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx);
-        assert!(cb.onSuccess(&mock_key).is_ok());
-
-        let key = tokio_rt().block_on(rx).unwrap().unwrap();
-        assert_eq!(key, mock_key);
-    }
-
-    #[test]
-    fn test_get_key_cb_cancel() {
-        let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx);
-        assert!(cb.onCancel().is_ok());
-
-        let result = tokio_rt().block_on(rx).unwrap();
-        assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::RequestCancelled);
-    }
-
-    #[test]
-    fn test_get_key_cb_error() {
-        for get_key_error in GetKeyErrorCode::enum_values() {
-            let (tx, rx) = oneshot::channel();
-            let cb = GetKeyCallback::new_native_binder(tx);
-            assert!(cb.onError(get_key_error, "error").is_ok());
-
-            let result = tokio_rt().block_on(rx).unwrap();
-            assert_eq!(
-                result.unwrap_err().downcast::<Error>().unwrap(),
-                Error::GetKeyFailed(get_key_error),
-            );
-        }
-    }
-
-    #[test]
-    fn test_store_upgraded_cb_success() {
-        let (tx, rx) = oneshot::channel();
-        let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
-        assert!(cb.onSuccess().is_ok());
-
-        tokio_rt().block_on(rx).unwrap().unwrap();
-    }
-
-    #[test]
-    fn test_store_upgraded_key_cb_error() {
-        let (tx, rx) = oneshot::channel();
-        let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
-        assert!(cb.onError("oh no! it failed").is_ok());
-
-        let result = tokio_rt().block_on(rx).unwrap();
-        assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::StoreUpgradedKeyFailed);
-    }
-
-    #[test]
-    fn test_get_mock_key_success() {
-        let mock_key =
-            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
-        let registration = get_mock_registration(&mock_key, /*latency=*/ None).unwrap();
-
-        let key = tokio_rt()
-            .block_on(get_rkpd_attestation_key_from_registration_async(&registration, 0))
-            .unwrap();
-        assert_eq!(key, mock_key);
-    }
-
-    #[test]
-    fn test_get_mock_key_timeout() {
-        let mock_key =
-            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
-        let latency = RKPD_TIMEOUT + Duration::from_secs(1);
-        let registration = get_mock_registration(&mock_key, Some(latency)).unwrap();
-
-        let result =
-            tokio_rt().block_on(get_rkpd_attestation_key_from_registration_async(&registration, 0));
-        assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::RetryableTimeout);
-    }
-
-    #[test]
-    fn test_store_mock_key_success() {
-        let mock_key =
-            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
-        let registration = get_mock_registration(&mock_key, /*latency=*/ None).unwrap();
-        tokio_rt()
-            .block_on(store_rkpd_attestation_key_with_registration_async(&registration, &[], &[]))
-            .unwrap();
-    }
-
-    #[test]
-    fn test_store_mock_key_timeout() {
-        let mock_key =
-            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
-        let latency = RKPD_TIMEOUT + Duration::from_secs(1);
-        let registration = get_mock_registration(&mock_key, Some(latency)).unwrap();
-
-        let result = tokio_rt().block_on(store_rkpd_attestation_key_with_registration_async(
-            &registration,
-            &[],
-            &[],
-        ));
-        assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::Timeout);
-    }
-
-    #[test]
-    fn test_get_rkpd_attestation_key() {
-        binder::ProcessState::start_thread_pool();
-        let key_id = get_next_key_id();
-        let key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
-        assert!(!key.keyBlob.is_empty());
-        assert!(!key.encodedCertChain.is_empty());
-    }
-
-    #[test]
-    fn test_get_rkpd_attestation_key_same_caller() {
-        binder::ProcessState::start_thread_pool();
-        let key_id = get_next_key_id();
-
-        // Multiple calls should return the same key.
-        let first_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
-        let second_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
-
-        assert_eq!(first_key.keyBlob, second_key.keyBlob);
-        assert_eq!(first_key.encodedCertChain, second_key.encodedCertChain);
-    }
-
-    #[test]
-    fn test_get_rkpd_attestation_key_different_caller() {
-        binder::ProcessState::start_thread_pool();
-        let first_key_id = get_next_key_id();
-        let second_key_id = get_next_key_id();
-
-        // Different callers should be getting different keys.
-        let first_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, first_key_id).unwrap();
-        let second_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, second_key_id).unwrap();
-
-        assert_ne!(first_key.keyBlob, second_key.keyBlob);
-        assert_ne!(first_key.encodedCertChain, second_key.encodedCertChain);
-    }
-
-    #[test]
-    // Couple of things to note:
-    // 1. This test must never run with UID of keystore. Otherwise, it can mess up keys stored by
-    //    keystore.
-    // 2. Storing and reading the stored key is prone to race condition. So, we only do this in one
-    //    test case.
-    fn test_store_rkpd_attestation_key() {
-        binder::ProcessState::start_thread_pool();
-        let key_id = get_next_key_id();
-        let key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
-        let new_blob: [u8; 8] = rand::random();
-
-        assert!(
-            store_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, &key.keyBlob, &new_blob).is_ok()
-        );
-
-        let new_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
-
-        // Restore original key so that we don't leave RKPD with invalid blobs.
-        assert!(
-            store_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, &new_blob, &key.keyBlob).is_ok()
-        );
-        assert_eq!(new_key.keyBlob, new_blob);
-    }
-
-    #[test]
-    fn test_stress_get_rkpd_attestation_key() {
-        binder::ProcessState::start_thread_pool();
-        let key_id = get_next_key_id();
-        let mut threads = vec![];
-        const NTHREADS: u32 = 10;
-        const NCALLS: u32 = 1000;
-
-        for _ in 0..NTHREADS {
-            threads.push(std::thread::spawn(move || {
-                for _ in 0..NCALLS {
-                    let key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
-                    assert!(!key.keyBlob.is_empty());
-                    assert!(!key.encodedCertChain.is_empty());
-                }
-            }));
-        }
-
-        for t in threads {
-            assert!(t.join().is_ok());
-        }
-    }
-}
+mod tests;
diff --git a/keystore2/rkpd_client/src/tests.rs b/keystore2/rkpd_client/src/tests.rs
new file mode 100644
index 0000000..fd0468f
--- /dev/null
+++ b/keystore2/rkpd_client/src/tests.rs
@@ -0,0 +1,338 @@
+// Copyright 2022, 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.
+
+//! RKPD tests.
+
+use super::*;
+use android_security_rkp_aidl::aidl::android::security::rkp::IRegistration::BnRegistration;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::{Arc, Mutex};
+
+const DEFAULT_RPC_SERVICE_NAME: &str =
+    "android.hardware.security.keymint.IRemotelyProvisionedComponent/default";
+
+struct MockRegistrationValues {
+    key: RemotelyProvisionedKey,
+    latency: Option<Duration>,
+    thread_join_handles: Vec<Option<std::thread::JoinHandle<()>>>,
+}
+
+struct MockRegistration(Arc<Mutex<MockRegistrationValues>>);
+
+impl MockRegistration {
+    pub fn new_native_binder(
+        key: &RemotelyProvisionedKey,
+        latency: Option<Duration>,
+    ) -> Strong<dyn IRegistration> {
+        let result = Self(Arc::new(Mutex::new(MockRegistrationValues {
+            key: RemotelyProvisionedKey {
+                keyBlob: key.keyBlob.clone(),
+                encodedCertChain: key.encodedCertChain.clone(),
+            },
+            latency,
+            thread_join_handles: Vec::new(),
+        })));
+        BnRegistration::new_binder(result, BinderFeatures::default())
+    }
+}
+
+impl Drop for MockRegistration {
+    fn drop(&mut self) {
+        let mut values = self.0.lock().unwrap();
+        for handle in values.thread_join_handles.iter_mut() {
+            // These are test threads. So, no need to worry too much about error handling.
+            handle.take().unwrap().join().unwrap();
+        }
+    }
+}
+
+impl Interface for MockRegistration {}
+
+impl IRegistration for MockRegistration {
+    fn getKey(&self, _: i32, cb: &Strong<dyn IGetKeyCallback>) -> binder::Result<()> {
+        let mut values = self.0.lock().unwrap();
+        let key = RemotelyProvisionedKey {
+            keyBlob: values.key.keyBlob.clone(),
+            encodedCertChain: values.key.encodedCertChain.clone(),
+        };
+        let latency = values.latency;
+        let get_key_cb = cb.clone();
+
+        // Need a separate thread to trigger timeout in the caller.
+        let join_handle = std::thread::spawn(move || {
+            if let Some(duration) = latency {
+                std::thread::sleep(duration);
+            }
+            get_key_cb.onSuccess(&key).unwrap();
+        });
+        values.thread_join_handles.push(Some(join_handle));
+        Ok(())
+    }
+
+    fn cancelGetKey(&self, _: &Strong<dyn IGetKeyCallback>) -> binder::Result<()> {
+        Ok(())
+    }
+
+    fn storeUpgradedKeyAsync(
+        &self,
+        _: &[u8],
+        _: &[u8],
+        cb: &Strong<dyn IStoreUpgradedKeyCallback>,
+    ) -> binder::Result<()> {
+        // We are primarily concerned with timing out correctly. Storing the key in this mock
+        // registration isn't particularly interesting, so skip that part.
+        let values = self.0.lock().unwrap();
+        let store_cb = cb.clone();
+        let latency = values.latency;
+
+        std::thread::spawn(move || {
+            if let Some(duration) = latency {
+                std::thread::sleep(duration);
+            }
+            store_cb.onSuccess().unwrap();
+        });
+        Ok(())
+    }
+}
+
+fn get_mock_registration(
+    key: &RemotelyProvisionedKey,
+    latency: Option<Duration>,
+) -> Result<binder::Strong<dyn IRegistration>> {
+    let (tx, rx) = oneshot::channel();
+    let cb = GetRegistrationCallback::new_native_binder(tx);
+    let mock_registration = MockRegistration::new_native_binder(key, latency);
+
+    assert!(cb.onSuccess(&mock_registration).is_ok());
+    tokio_rt().block_on(rx).unwrap()
+}
+
+// Using the same key ID makes test cases race with each other. So, we use separate key IDs for
+// different test cases.
+fn get_next_key_id() -> u32 {
+    static ID: AtomicU32 = AtomicU32::new(0);
+    ID.fetch_add(1, Ordering::Relaxed)
+}
+
+#[test]
+fn test_get_registration_cb_success() {
+    let key: RemotelyProvisionedKey = Default::default();
+    let registration = get_mock_registration(&key, /*latency=*/ None);
+    assert!(registration.is_ok());
+}
+
+#[test]
+fn test_get_registration_cb_cancel() {
+    let (tx, rx) = oneshot::channel();
+    let cb = GetRegistrationCallback::new_native_binder(tx);
+    assert!(cb.onCancel().is_ok());
+
+    let result = tokio_rt().block_on(rx).unwrap();
+    assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::RequestCancelled);
+}
+
+#[test]
+fn test_get_registration_cb_error() {
+    let (tx, rx) = oneshot::channel();
+    let cb = GetRegistrationCallback::new_native_binder(tx);
+    assert!(cb.onError("error").is_ok());
+
+    let result = tokio_rt().block_on(rx).unwrap();
+    assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::GetRegistrationFailed);
+}
+
+#[test]
+fn test_get_key_cb_success() {
+    let mock_key =
+        RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+    let (tx, rx) = oneshot::channel();
+    let cb = GetKeyCallback::new_native_binder(tx);
+    assert!(cb.onSuccess(&mock_key).is_ok());
+
+    let key = tokio_rt().block_on(rx).unwrap().unwrap();
+    assert_eq!(key, mock_key);
+}
+
+#[test]
+fn test_get_key_cb_cancel() {
+    let (tx, rx) = oneshot::channel();
+    let cb = GetKeyCallback::new_native_binder(tx);
+    assert!(cb.onCancel().is_ok());
+
+    let result = tokio_rt().block_on(rx).unwrap();
+    assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::RequestCancelled);
+}
+
+#[test]
+fn test_get_key_cb_error() {
+    for get_key_error in GetKeyErrorCode::enum_values() {
+        let (tx, rx) = oneshot::channel();
+        let cb = GetKeyCallback::new_native_binder(tx);
+        assert!(cb.onError(get_key_error, "error").is_ok());
+
+        let result = tokio_rt().block_on(rx).unwrap();
+        assert_eq!(
+            result.unwrap_err().downcast::<Error>().unwrap(),
+            Error::GetKeyFailed(get_key_error),
+        );
+    }
+}
+
+#[test]
+fn test_store_upgraded_cb_success() {
+    let (tx, rx) = oneshot::channel();
+    let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+    assert!(cb.onSuccess().is_ok());
+
+    tokio_rt().block_on(rx).unwrap().unwrap();
+}
+
+#[test]
+fn test_store_upgraded_key_cb_error() {
+    let (tx, rx) = oneshot::channel();
+    let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+    assert!(cb.onError("oh no! it failed").is_ok());
+
+    let result = tokio_rt().block_on(rx).unwrap();
+    assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::StoreUpgradedKeyFailed);
+}
+
+#[test]
+fn test_get_mock_key_success() {
+    let mock_key =
+        RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+    let registration = get_mock_registration(&mock_key, /*latency=*/ None).unwrap();
+
+    let key = tokio_rt()
+        .block_on(get_rkpd_attestation_key_from_registration_async(&registration, 0))
+        .unwrap();
+    assert_eq!(key, mock_key);
+}
+
+#[test]
+fn test_get_mock_key_timeout() {
+    let mock_key =
+        RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+    let latency = RKPD_TIMEOUT + Duration::from_secs(1);
+    let registration = get_mock_registration(&mock_key, Some(latency)).unwrap();
+
+    let result =
+        tokio_rt().block_on(get_rkpd_attestation_key_from_registration_async(&registration, 0));
+    assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::RetryableTimeout);
+}
+
+#[test]
+fn test_store_mock_key_success() {
+    let mock_key =
+        RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+    let registration = get_mock_registration(&mock_key, /*latency=*/ None).unwrap();
+    tokio_rt()
+        .block_on(store_rkpd_attestation_key_with_registration_async(&registration, &[], &[]))
+        .unwrap();
+}
+
+#[test]
+fn test_store_mock_key_timeout() {
+    let mock_key =
+        RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+    let latency = RKPD_TIMEOUT + Duration::from_secs(1);
+    let registration = get_mock_registration(&mock_key, Some(latency)).unwrap();
+
+    let result = tokio_rt().block_on(store_rkpd_attestation_key_with_registration_async(
+        &registration,
+        &[],
+        &[],
+    ));
+    assert_eq!(result.unwrap_err().downcast::<Error>().unwrap(), Error::Timeout);
+}
+
+#[test]
+fn test_get_rkpd_attestation_key() {
+    binder::ProcessState::start_thread_pool();
+    let key_id = get_next_key_id();
+    let key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
+    assert!(!key.keyBlob.is_empty());
+    assert!(!key.encodedCertChain.is_empty());
+}
+
+#[test]
+fn test_get_rkpd_attestation_key_same_caller() {
+    binder::ProcessState::start_thread_pool();
+    let key_id = get_next_key_id();
+
+    // Multiple calls should return the same key.
+    let first_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
+    let second_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
+
+    assert_eq!(first_key.keyBlob, second_key.keyBlob);
+    assert_eq!(first_key.encodedCertChain, second_key.encodedCertChain);
+}
+
+#[test]
+fn test_get_rkpd_attestation_key_different_caller() {
+    binder::ProcessState::start_thread_pool();
+    let first_key_id = get_next_key_id();
+    let second_key_id = get_next_key_id();
+
+    // Different callers should be getting different keys.
+    let first_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, first_key_id).unwrap();
+    let second_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, second_key_id).unwrap();
+
+    assert_ne!(first_key.keyBlob, second_key.keyBlob);
+    assert_ne!(first_key.encodedCertChain, second_key.encodedCertChain);
+}
+
+#[test]
+// Couple of things to note:
+// 1. This test must never run with UID of keystore. Otherwise, it can mess up keys stored by
+//    keystore.
+// 2. Storing and reading the stored key is prone to race condition. So, we only do this in one
+//    test case.
+fn test_store_rkpd_attestation_key() {
+    binder::ProcessState::start_thread_pool();
+    let key_id = get_next_key_id();
+    let key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
+    let new_blob: [u8; 8] = rand::random();
+
+    assert!(store_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, &key.keyBlob, &new_blob).is_ok());
+
+    let new_key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
+
+    // Restore original key so that we don't leave RKPD with invalid blobs.
+    assert!(store_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, &new_blob, &key.keyBlob).is_ok());
+    assert_eq!(new_key.keyBlob, new_blob);
+}
+
+#[test]
+fn test_stress_get_rkpd_attestation_key() {
+    binder::ProcessState::start_thread_pool();
+    let key_id = get_next_key_id();
+    let mut threads = vec![];
+    const NTHREADS: u32 = 10;
+    const NCALLS: u32 = 1000;
+
+    for _ in 0..NTHREADS {
+        threads.push(std::thread::spawn(move || {
+            for _ in 0..NCALLS {
+                let key = get_rkpd_attestation_key(DEFAULT_RPC_SERVICE_NAME, key_id).unwrap();
+                assert!(!key.keyBlob.is_empty());
+                assert!(!key.encodedCertChain.is_empty());
+            }
+        }));
+    }
+
+    for t in threads {
+        assert!(t.join().is_ok());
+    }
+}