Address comments

src/env/test/mod.rs

@@ -36,6 +36,7 @@ pub struct TestEnv {
     store: Store<BufferStorage>,
     upgrade_storage: Option<BufferUpgradeStorage>,
     customization: TestCustomization,
+    attestation_id: attestation_store::Id,
 }
 
 pub struct TestRng256 {
@@ -106,12 +107,14 @@ impl TestEnv {
         let store = Store::new(storage).ok().unwrap();
         let upgrade_storage = Some(BufferUpgradeStorage::new().unwrap());
         let customization = DEFAULT_CUSTOMIZATION.into();
+        let attestation_id = attestation_store::Id::Batch;
         TestEnv {
             rng,
             user_presence,
             store,
             upgrade_storage,
             customization,
+            attestation_id,
         }
     }
 
@@ -126,6 +129,10 @@ impl TestEnv {
     pub fn rng(&mut self) -> &mut TestRng256 {
         &mut self.rng
     }
+
+    pub fn set_attestation_id(&mut self, id: attestation_store::Id) {
+        self.attestation_id = id;
+    }
 }
 
 impl TestUserPresence {
@@ -149,7 +156,12 @@ impl FirmwareProtection for TestEnv {
 }
 
 impl key_store::Helper for TestEnv {}
-impl attestation_store::Helper for TestEnv {}
+
+impl attestation_store::Helper for TestEnv {
+    fn attestation_id(&self) -> attestation_store::Id {
+        self.attestation_id.clone()
+    }
+}
 
 impl Env for TestEnv {
     type Rng = TestRng256;