diff --git a/keystore2/test_utils/authorizations.rs b/keystore2/test_utils/authorizations.rs
index aa75982..b73aab5 100644
--- a/keystore2/test_utils/authorizations.rs
+++ b/keystore2/test_utils/authorizations.rs
@@ -269,6 +269,42 @@
         });
         self
     }
+
+    /// Set boot loader only.
+    pub fn boot_loader_only(mut self) -> Self {
+        self.0.push(KeyParameter {
+            tag: Tag::BOOTLOADER_ONLY,
+            value: KeyParameterValue::BoolValue(true),
+        });
+        self
+    }
+
+    /// Set early boot only.
+    pub fn early_boot_only(mut self) -> Self {
+        self.0.push(KeyParameter {
+            tag: Tag::EARLY_BOOT_ONLY,
+            value: KeyParameterValue::BoolValue(true),
+        });
+        self
+    }
+
+    /// Set max uses per boot.
+    pub fn max_uses_per_boot(mut self, max_uses: i32) -> Self {
+        self.0.push(KeyParameter {
+            tag: Tag::MAX_USES_PER_BOOT,
+            value: KeyParameterValue::Integer(max_uses),
+        });
+        self
+    }
+
+    /// Set max usage count.
+    pub fn usage_count_limit(mut self, usage_count: i32) -> Self {
+        self.0.push(KeyParameter {
+            tag: Tag::USAGE_COUNT_LIMIT,
+            value: KeyParameterValue::Integer(usage_count),
+        });
+        self
+    }
 }
 
 impl Deref for AuthSetBuilder {
diff --git a/keystore2/test_utils/ffi_test_utils.cpp b/keystore2/test_utils/ffi_test_utils.cpp
index 0740804..4e781d1 100644
--- a/keystore2/test_utils/ffi_test_utils.cpp
+++ b/keystore2/test_utils/ffi_test_utils.cpp
@@ -601,7 +601,8 @@
     return result;
 }
 
-CxxResult getValueFromAttestRecord(rust::Vec<rust::u8> cert_buf, int32_t tag) {
+CxxResult getValueFromAttestRecord(rust::Vec<rust::u8> cert_buf, int32_t tag,
+                                   int32_t expected_sec_level) {
     CxxResult cxx_result{};
     cxx_result.error = false;
 
@@ -649,6 +650,8 @@
 
     aidl::android::hardware::security::keymint::Tag auth_tag =
         static_cast<aidl::android::hardware::security::keymint::Tag>(tag);
+    aidl::android::hardware::security::keymint::SecurityLevel tag_security_level =
+        static_cast<aidl::android::hardware::security::keymint::SecurityLevel>(expected_sec_level);
 
     if (auth_tag == aidl::android::hardware::security::keymint::Tag::ATTESTATION_APPLICATION_ID) {
         int pos = att_sw_enforced.find(
@@ -685,6 +688,36 @@
         return cxx_result;
     }
 
+    if (auth_tag == aidl::android::hardware::security::keymint::Tag::USAGE_COUNT_LIMIT) {
+        aidl::android::hardware::security::keymint::KeyParameter param;
+        int pos = att_hw_enforced.find(auth_tag);
+        if (tag_security_level ==
+                aidl::android::hardware::security::keymint::SecurityLevel::SOFTWARE ||
+            tag_security_level ==
+                aidl::android::hardware::security::keymint::SecurityLevel::KEYSTORE) {
+            pos = att_sw_enforced.find(auth_tag);
+            if (pos == -1) {
+                LOG(ERROR) << "USAGE_COUNT_LIMIT not found in software enforced auth list";
+                cxx_result.error = KM_ERROR_INVALID_TAG;
+                return cxx_result;
+            }
+            param = att_sw_enforced[pos];
+        } else {
+            pos = att_hw_enforced.find(auth_tag);
+            if (pos == -1) {
+                LOG(ERROR) << "USAGE_COUNT_LIMIT not found in hardware enforced auth list";
+                cxx_result.error = KM_ERROR_INVALID_TAG;
+                return cxx_result;
+            }
+            param = att_hw_enforced[pos];
+        }
+        std::string val = std::to_string(
+            param.value
+                .get<aidl::android::hardware::security::keymint::KeyParameterValue::integer>());
+        std::move(val.begin(), val.end(), std::back_inserter(cxx_result.data));
+        return cxx_result;
+    }
+
     int pos = att_hw_enforced.find(auth_tag);
     if (pos == -1) {
         LOG(ERROR) << "getValueFromAttestRecord - unsupported tag.";
diff --git a/keystore2/test_utils/ffi_test_utils.hpp b/keystore2/test_utils/ffi_test_utils.hpp
index 69f558a..c4db1ba 100644
--- a/keystore2/test_utils/ffi_test_utils.hpp
+++ b/keystore2/test_utils/ffi_test_utils.hpp
@@ -1,16 +1,16 @@
 #pragma once
 
-#include "rust/cxx.h"
 #include "ffi_test_utils.rs.h"
+#include "rust/cxx.h"
 
 bool validateCertChain(rust::Vec<rust::u8> cert_buf, uint32_t cert_len, bool strict_issuer_check);
 CxxResult createWrappedKey(rust::Vec<rust::u8> encrypted_secure_key,
-                              rust::Vec<rust::u8> encrypted_transport_key,
-                              rust::Vec<rust::u8> iv,
-                              rust::Vec<rust::u8> tag);
+                           rust::Vec<rust::u8> encrypted_transport_key, rust::Vec<rust::u8> iv,
+                           rust::Vec<rust::u8> tag);
 CxxResult buildAsn1DerEncodedWrappedKeyDescription();
 bool performCryptoOpUsingKeystoreEngine(int64_t grant_id);
-CxxResult getValueFromAttestRecord(rust::Vec<rust::u8> cert_buf, int32_t tag);
+CxxResult getValueFromAttestRecord(rust::Vec<rust::u8> cert_buf, int32_t tag,
+                                   int32_t expected_sec_level);
 uint32_t getOsVersion();
 uint32_t getOsPatchlevel();
 uint32_t getVendorPatchlevel();
diff --git a/keystore2/test_utils/ffi_test_utils.rs b/keystore2/test_utils/ffi_test_utils.rs
index 04d82f1..5d6bf46 100644
--- a/keystore2/test_utils/ffi_test_utils.rs
+++ b/keystore2/test_utils/ffi_test_utils.rs
@@ -15,7 +15,9 @@
 //! This module implements helper methods to access the functionalities implemented in CPP.
 
 use crate::key_generations::Error;
-use android_hardware_security_keymint::aidl::android::hardware::security::keymint::Tag::Tag;
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    SecurityLevel::SecurityLevel, Tag::Tag,
+};
 
 #[cxx::bridge]
 mod ffi {
@@ -35,7 +37,11 @@
         ) -> CxxResult;
         fn buildAsn1DerEncodedWrappedKeyDescription() -> CxxResult;
         fn performCryptoOpUsingKeystoreEngine(grant_id: i64) -> bool;
-        fn getValueFromAttestRecord(cert_buf: Vec<u8>, tag: i32) -> CxxResult;
+        fn getValueFromAttestRecord(
+            cert_buf: Vec<u8>,
+            tag: i32,
+            expected_sec_level: i32,
+        ) -> CxxResult;
         fn getOsVersion() -> u32;
         fn getOsPatchlevel() -> u32;
         fn getVendorPatchlevel() -> u32;
@@ -98,8 +104,12 @@
 }
 
 /// Get the value of the given `Tag` from attestation record.
-pub fn get_value_from_attest_record(cert_buf: &[u8], tag: Tag) -> Result<Vec<u8>, Error> {
-    let result = ffi::getValueFromAttestRecord(cert_buf.to_vec(), tag.0);
+pub fn get_value_from_attest_record(
+    cert_buf: &[u8],
+    tag: Tag,
+    expected_sec_level: SecurityLevel,
+) -> Result<Vec<u8>, Error> {
+    let result = ffi::getValueFromAttestRecord(cert_buf.to_vec(), tag.0, expected_sec_level.0);
     if !result.error && !result.data.is_empty() {
         return Ok(result.data);
     }
diff --git a/keystore2/test_utils/key_generations.rs b/keystore2/test_utils/key_generations.rs
index 24ce6e1..ccf27bc 100644
--- a/keystore2/test_utils/key_generations.rs
+++ b/keystore2/test_utils/key_generations.rs
@@ -470,6 +470,18 @@
     ));
 }
 
+/// Get the key `Authorization` for the given auth `Tag`.
+pub fn get_key_auth(authorizations: &[Authorization], tag: Tag) -> Option<&Authorization> {
+    let auths: Vec<&Authorization> =
+        authorizations.iter().filter(|auth| auth.keyParameter.tag == tag).collect();
+
+    if !auths.is_empty() {
+        Some(auths[0])
+    } else {
+        None
+    }
+}
+
 /// Generate EC Key using given security level and domain with below key parameters and
 /// optionally allow the generated key to be attested with factory provisioned attest key using
 /// given challenge and application id -