From 10d17f83746c52a211aea11549ee5245fe5437bf Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Fri, 4 Apr 2025 13:32:33 -0700
Subject: [PATCH 01/13] chore(dafny): KSA Create Branch Key with MRK keys

---
 .../test/TestCreateKeys.dfy                   | 109 ++++++++++--------
 .../test/TestGetKeys.dfy                      |  20 +++-
 .../test/AdminFixtures.dfy                    |   3 +-
 .../test/Mutations/TestDecryptEncrypt.dfy     |   1 +
 .../test/TestAdminHV1Only.dfy                 |  43 +++++++
 5 files changed, 127 insertions(+), 49 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
index 9c6e2f99ed..ff46d8e727 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
@@ -23,48 +23,63 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
   import AwsArnParsing
   import TestGetKeys
 
-  /*
-    // If you need to re-create the MRK Branch Keys
-  
-    method {:test} TestCreateMRK()
-    {
-      var ddbClient :- expect DDB.DynamoDBClient();
-  
-      var keyStoreConfigEast := Types.KeyStoreConfig(
-        id := None,
-        kmsConfiguration := KmsConfigEast,
-        logicalKeyStoreName := logicalKeyStoreName,
-        grantTokens := None,
-        ddbTableName := branchKeyStoreName,
-        ddbClient := Some(ddbClient)
-      );
-  
-      var keyStoreConfigWest := Types.KeyStoreConfig(
-        id := None,
-        kmsConfiguration := KmsConfigWest,
-        logicalKeyStoreName := logicalKeyStoreName,
-        grantTokens := None,
-        ddbTableName := branchKeyStoreName,
-        ddbClient := Some(ddbClient)
-      );
-  
-      var keyStoreEast :- expect KeyStore.KeyStore(keyStoreConfigEast);
-      var keyStoreWest :- expect KeyStore.KeyStore(keyStoreConfigWest);
-  
-      var branchKeyIdWest :- expect keyStoreWest.CreateKey(Types.CreateKeyInput(
-                                                             branchKeyIdentifier := Some(WestBranchKey),
-                                                             encryptionContext := Some(KmsMrkEC)
-                                                           ));
-  
-      var branchKeyIdEast :- expect keyStoreEast.CreateKey(Types.CreateKeyInput(
-                                                             branchKeyIdentifier := Some(EastBranchKey),
-                                                             encryptionContext := Some(KmsMrkEC)
-                                                           ));
-  
-      expect branchKeyIdEast.branchKeyIdentifier == EastBranchKey;
-      expect branchKeyIdWest.branchKeyIdentifier == WestBranchKey;
-    }
-  */
+  method {:test} TestCreateMRKForHV1()
+  {
+    var ddbClient :- expect DDB.DynamoDBClient();
+
+    var keyStoreConfigEast := Types.KeyStoreConfig(
+      id := None,
+      kmsConfiguration := KmsConfigEast,
+      logicalKeyStoreName := logicalKeyStoreName,
+      grantTokens := None,
+      storage := Some(
+        Types.ddb(
+          Types.DynamoDBTable(
+            ddbTableName := branchKeyStoreName,
+            ddbClient := Some(ddbClient)
+          )))
+    );
+
+    var keyStoreConfigWest := Types.KeyStoreConfig(
+      id := None,
+      kmsConfiguration := KmsConfigWest,
+      logicalKeyStoreName := logicalKeyStoreName,
+      grantTokens := None,
+      storage := Some(
+        Types.ddb(
+          Types.DynamoDBTable(
+            ddbTableName := branchKeyStoreName,
+            ddbClient := Some(ddbClient)
+          )))
+    );
+
+    // Create key with Custom EC & Branch Key Identifier
+    var uuid :- expect UUID.GenerateUUID();
+    var happyCaseId := "test-happy-case-create-key-hv-1";
+    var branchKeyIdWest := happyCaseId + "-" + "west" + "-" + uuid;
+    // print branchKeyIdWest;
+    var branchKeyIdEast := happyCaseId + "-" + "east" + "-" + uuid;
+    // print branchKeyIdEast;
+
+    var keyStoreEast :- expect KeyStore.KeyStore(keyStoreConfigEast);
+    var keyStoreWest :- expect KeyStore.KeyStore(keyStoreConfigWest);
+
+    var branchKeyIdWest? :- expect keyStoreWest.CreateKey(Types.CreateKeyInput(
+                                                            branchKeyIdentifier := Some(branchKeyIdWest),
+                                                            encryptionContext := Some(KmsMrkEC)
+                                                          ));
+
+    var branchKeyIdEast? :- expect keyStoreEast.CreateKey(Types.CreateKeyInput(
+                                                            branchKeyIdentifier := Some(branchKeyIdEast),
+                                                            encryptionContext := Some(KmsMrkEC)
+                                                          ));
+
+    expect branchKeyIdEast?.branchKeyIdentifier == branchKeyIdEast;
+    expect branchKeyIdWest?.branchKeyIdentifier == branchKeyIdWest;
+
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyIdWest, ddbClient:=ddbClient);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyIdEast, ddbClient:=ddbClient);
+  }
 
   method {:test} TestCreateBranchAndBeaconKeys()
   {
@@ -81,12 +96,12 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
           Types.DynamoDBTable(
             ddbTableName := branchKeyStoreName,
             ddbClient := Some(ddbClient)
-          ))),
-      keyManagement := Some(
-        Types.kms(
-          Types.AwsKms(
-            kmsClient := Some(kmsClient)
           )))
+      // keyManagement := Some(
+      //   Types.kms(
+      //     Types.AwsKms(
+      //       kmsClient := Some(kmsClient)
+      //     )))
     );
 
     var keyStore :- expect KeyStore.KeyStore(keyStoreConfig);
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
index 2309503b51..2eb02d8700 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
@@ -72,11 +72,29 @@ module TestGetKeys {
     testBranchKeyVersionHappyCase(keyStore, hv2BranchKeyId, hv2BranchKeyVersion, hv2BranchKeyIdActiveVersionUtf8Bytes);
   }
 
+  // This is a static test case to Get Branch Keys created with Mrk Keys
+  method {:test} TestGetActiveMrkKey()
+  {
+    VerifyGetActiveMrkKey(
+      KmsConfigEast := KmsConfigEast,
+      KmsConfigWest := KmsConfigWest,
+      KmsMrkConfigEast := KmsMrkConfigEast,
+      KmsMrkConfigWest := KmsMrkConfigWest
+    );
+  }
+
   // TODO-HV2-M1: Add MRK test for hv2
-  method {:test} {:isolate_assertions} TestGetActiveMrkKey()
+  method {:isolate_assertions} VerifyGetActiveMrkKey(
+    KmsConfigEast : Types.KMSConfiguration,
+    KmsConfigWest : Types.KMSConfiguration,
+    KmsMrkConfigEast : Types.KMSConfiguration,
+    KmsMrkConfigWest : Types.KMSConfiguration
+  )
   {
     var ddbClient :- expect DDB.DynamoDBClient();
 
+    print "Running MRK";
+
     var eastKeyStoreConfig := Types.KeyStoreConfig(
       id := None,
       kmsConfiguration := KmsConfigEast,
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/AdminFixtures.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/AdminFixtures.dfy
index da3df95ec5..a4fc543fb7 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/AdminFixtures.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/AdminFixtures.dfy
@@ -44,7 +44,8 @@ module {:options "/functionSyntax:4" } AdminFixtures {
       ddbClient := ddbClient,
       logicalKeyStoreName := logicalName,
       ddbTableNameUtf8 := physicalNameUtf8,
-      logicalKeyStoreNameUtf8 := logicalNameUtf8);
+      logicalKeyStoreNameUtf8 := logicalNameUtf8
+    );
 
     var underTestConfig := Types.KeyStoreAdminConfig(
       logicalKeyStoreName := logicalName,
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/Mutations/TestDecryptEncrypt.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/Mutations/TestDecryptEncrypt.dfy
index 0fd671557e..9bdf790430 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/Mutations/TestDecryptEncrypt.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/Mutations/TestDecryptEncrypt.dfy
@@ -82,6 +82,7 @@ module {:options "/functionSyntax:4" } TestDecryptEncryptStrat {
 
     var timestamp :- expect Time.GetCurrentTimeStamp();
     var newCustomEC: KeyStoreTypes.EncryptionContextString := map["Koda" := timestamp];
+    // TODO-HV-2-M2: We should Mutate with Mrk Key as well
     var mutationsRequest := Types.Mutations(TerminalEncryptionContext := Some(newCustomEC));
     var initInput := Types.InitializeMutationInput(
       Identifier := testId,
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
index 21db35c333..e761ad8f83 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
@@ -100,6 +100,49 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyId, ddbClient:=ddbClient);
   }
 
+  // TODO-HV-2-M2?: Document that BKSA ONLY calls KMS/DDB for the region of the supplied clients.
+  // or Ensure BKSA has similar behavior as BKS in terms of MRK region.
+  /*
+  method {:test} TestCreateMRKForHV2()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+
+    // Create key with Custom EC & Branch Key Identifier
+    var uuid :- expect UUID.GenerateUUID();
+    var branchKeyIdWest := happyCaseId + "-" + WestBranchKey + "-" + uuid;
+    var branchKeyIdEast := happyCaseId + "-" + EastBranchKey + "-" + uuid;
+
+    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
+
+    var westOutput? :- expect underTest.CreateKey(Types.CreateKeyInput(
+                                                    Identifier := Some(branchKeyIdWest),
+                                                    EncryptionContext := Some(customEC),
+                                                    KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnWest),
+                                                    Strategy := Some(strategy),
+                                                    HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+                                                  ));
+    expect westOutput?.Identifier == branchKeyIdWest;
+    expect westOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
+    print branchKeyIdWest + "\n";
+
+    var eastOutput? :- expect underTest.CreateKey(Types.CreateKeyInput(
+                                                    Identifier := Some(branchKeyIdEast),
+                                                    EncryptionContext := Some(customEC),
+                                                    KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnEast),
+                                                    Strategy := Some(strategy),
+                                                    HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+                                                  ));
+    expect eastOutput?.Identifier == branchKeyIdEast;
+    expect eastOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
+    print branchKeyIdEast + "\n";
+  }
+  */
+
   // TODO-HV-2-M2 : Probably make this a happy test?
   const testMutateForHV2FailsCaseId := "dafny-initialize-mutation-hv-2-rejection"
   method {:test} TestMutateForHV2Fails()

From 1ac074a690202e97ac7a58dddc9cd08fbd3c0942 Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Fri, 4 Apr 2025 13:43:13 -0700
Subject: [PATCH 02/13] nit

---
 .../dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
index ff46d8e727..e3abddc022 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
@@ -23,6 +23,8 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
   import AwsArnParsing
   import TestGetKeys
 
+  const happyCaseId := "test-happy-case-create-key-hv-1"
+
   method {:test} TestCreateMRKForHV1()
   {
     var ddbClient :- expect DDB.DynamoDBClient();
@@ -55,7 +57,6 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
 
     // Create key with Custom EC & Branch Key Identifier
     var uuid :- expect UUID.GenerateUUID();
-    var happyCaseId := "test-happy-case-create-key-hv-1";
     var branchKeyIdWest := happyCaseId + "-" + "west" + "-" + uuid;
     // print branchKeyIdWest;
     var branchKeyIdEast := happyCaseId + "-" + "east" + "-" + uuid;

From 89cbfc354de3c42fc3e60ec569b50b29ef5e4da9 Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Fri, 4 Apr 2025 15:46:32 -0700
Subject: [PATCH 03/13] Add more tests

---
 ...{TestAdminHV1Only.dfy => TestAdminHV2.dfy} | 72 +++++++++++++++++--
 1 file changed, 66 insertions(+), 6 deletions(-)
 rename AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/{TestAdminHV1Only.dfy => TestAdminHV2.dfy} (81%)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
similarity index 81%
rename from AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
rename to AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
index e761ad8f83..c8821ed459 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
@@ -8,9 +8,10 @@ include "../../AwsCryptographyKeyStore/test/TestGetKeys.dfy"
 include "../../AwsCryptographyKeyStore/Model/AwsCryptographyKeyStoreTypes.dfy"
 include "AdminFixtures.dfy"
 
-// TODO-HV-2-M1 : remove HV-1 restriction for CreateKey
 // TODO-HV-2-M2 : remove HV-1 restriction for Mutations
-module {:options "/functionSyntax:4" } TestAdminHV1Only {
+// TODO-HV-2-M3 : rename TestAdminHV2 -> TestAdminCreateKeyHV2
+//    and move HappyCases for version & mutations to it's own test modules
+module {:options "/functionSyntax:4" } TestAdminHV2 {
   import Types = AwsCryptographyKeyStoreAdminTypes
   import KeyStoreAdmin
   import KeyStore
@@ -26,9 +27,7 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
   import AdminFixtures
   import TestGetKeys
 
-  const happyCaseId := "test-create-key-hv-2-happy-case"
-
-  method {:test} TestCreateKeyForHV2HappyCase()
+  method {:test} TestCreateKeyHV2HappyCase()
   {
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
@@ -61,7 +60,8 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     var _ := CleanupItems.DeleteBranchKey(Identifier:=createKeyOutput?.Identifier, ddbClient:=ddbClient);
   }
 
-  method {:test} TestCreateKeyForHV2CreateOptions() {
+  const happyCaseId := "test-create-key-hv2-happy-case"
+  method {:test} TestCreateKeyHV2WithOptions() {
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
     var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
@@ -142,6 +142,66 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     print branchKeyIdEast + "\n";
   }
   */
+  
+  const testCreateKeyWithBKIDandNoECFailsId := "create-key-hv2-bkid-no-ec-fail"
+  method {:test} TestCreateKeyBKIDNoECFails() 
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+
+    var uuid :- expect UUID.GenerateUUID();
+    var branchKeyId := testCreateKeyWithBKIDandNoECFailsId + "-" + uuid;
+
+    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
+
+    var input := Types.CreateKeyInput(
+      Identifier := Some(branchKeyId),
+      EncryptionContext := None,
+      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+      Strategy := Some(strategy),
+      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+    );
+
+    var createKeyOutput := underTest.CreateKey(input);
+    expect createKeyOutput.Failure?, "Expected a failure when provided with Branch Key Identifier with no Custom Encryption Context";
+  }
+
+  method {:test} TestCreateKeyCustomECOnly() 
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+
+    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
+
+    var input := Types.CreateKeyInput(
+      Identifier := None,
+      EncryptionContext := Some(customEC),
+      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+      Strategy := Some(strategy),
+      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+    );
+
+    var createKeyOutput? :- expect underTest.CreateKey(input);
+    expect createKeyOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
+
+    // Get branch key items from storage
+    TestGetKeys.VerifyGetKeys(
+      identifier := createKeyOutput?.Identifier,
+      keyStore := keyStore,
+      storage := storage
+    );
+
+    // Since this process uses a read DDB table,
+    // the number of records will forever increase.
+    // To avoid this, remove the items.
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=createKeyOutput?.Identifier, ddbClient:=ddbClient);
+  }
 
   // TODO-HV-2-M2 : Probably make this a happy test?
   const testMutateForHV2FailsCaseId := "dafny-initialize-mutation-hv-2-rejection"

From cd3c475556bad00b50fcc15954c900fb1d429eb3 Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Fri, 4 Apr 2025 15:54:38 -0700
Subject: [PATCH 04/13] format

---
 .../AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy      | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
index c8821ed459..7bb12f209e 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
@@ -142,9 +142,9 @@ module {:options "/functionSyntax:4" } TestAdminHV2 {
     print branchKeyIdEast + "\n";
   }
   */
-  
+
   const testCreateKeyWithBKIDandNoECFailsId := "create-key-hv2-bkid-no-ec-fail"
-  method {:test} TestCreateKeyBKIDNoECFails() 
+  method {:test} TestCreateKeyBKIDNoECFails()
   {
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
@@ -168,7 +168,7 @@ module {:options "/functionSyntax:4" } TestAdminHV2 {
     expect createKeyOutput.Failure?, "Expected a failure when provided with Branch Key Identifier with no Custom Encryption Context";
   }
 
-  method {:test} TestCreateKeyCustomECOnly() 
+  method {:test} TestCreateKeyCustomECOnly()
   {
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();

From b3612bdf3343e835e6d7577e28c2eee0abd0745e Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Fri, 4 Apr 2025 16:50:23 -0700
Subject: [PATCH 05/13] Apply suggestions from code review

Co-authored-by: Tony Knapp <5892063+texastony@users.noreply.github.com>
---
 .../AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy      | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
index 7bb12f209e..a2d49bcef8 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
@@ -153,9 +153,6 @@ module {:options "/functionSyntax:4" } TestAdminHV2 {
 
     var uuid :- expect UUID.GenerateUUID();
     var branchKeyId := testCreateKeyWithBKIDandNoECFailsId + "-" + uuid;
-
-    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
-
     var input := Types.CreateKeyInput(
       Identifier := Some(branchKeyId),
       EncryptionContext := None,
@@ -197,9 +194,6 @@ module {:options "/functionSyntax:4" } TestAdminHV2 {
       storage := storage
     );
 
-    // Since this process uses a read DDB table,
-    // the number of records will forever increase.
-    // To avoid this, remove the items.
     var _ := CleanupItems.DeleteBranchKey(Identifier:=createKeyOutput?.Identifier, ddbClient:=ddbClient);
   }
 

From f5312ca42e3187dbdef1ea12224722e5b39b763e Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Sun, 6 Apr 2025 02:26:49 -0700
Subject: [PATCH 06/13] feedback

---
 .../test/TestCreateKeys.dfy                   | 10 +--
 .../test/TestGetKeys.dfy                      | 87 +++++++++++++++----
 .../test/TestAdminHV2.dfy                     |  2 +
 3 files changed, 77 insertions(+), 22 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
index e3abddc022..500a6d4995 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
@@ -97,12 +97,12 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
           Types.DynamoDBTable(
             ddbTableName := branchKeyStoreName,
             ddbClient := Some(ddbClient)
+          ))),
+      keyManagement := Some(
+        Types.kms(
+          Types.AwsKms(
+            kmsClient := Some(kmsClient)
           )))
-      // keyManagement := Some(
-      //   Types.kms(
-      //     Types.AwsKms(
-      //       kmsClient := Some(kmsClient)
-      //     )))
     );
 
     var keyStore :- expect KeyStore.KeyStore(keyStoreConfig);
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
index 2eb02d8700..7860558c52 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
@@ -27,7 +27,8 @@ module TestGetKeys {
   method VerifyGetKeys(
     identifier : string,
     keyStore : Types.IKeyStoreClient,
-    storage : Types.IKeyStorageInterface
+    storage : Types.IKeyStorageInterface,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
   )
     requires
       keyStore.ValidState() && storage.ValidState()
@@ -37,12 +38,26 @@ module TestGetKeys {
       keyStore.ValidState() && storage.ValidState()
 
   {
-    testBeaconKeyHappyCase(keyStore, identifier);
+    testBeaconKeyHappyCase(
+      keyStore := keyStore,
+      branchKeyId := identifier,
+      encryptionContext := encryptionContext
+    );
 
-    var activeResult := testAndGetActiveBranchKeyHappyCase(keyStore, identifier);
+    var activeResult := testAndGetActiveBranchKeyHappyCase(
+      keyStore := keyStore,
+      branchKeyId := identifier,
+      encryptionContext := encryptionContext
+    );
     var branchKeyVersion :- expect UTF8.Decode(activeResult.branchKeyMaterials.branchKeyVersion);
 
-    testBranchKeyVersionHappyCase(keyStore, identifier, branchKeyVersion, activeResult.branchKeyMaterials.branchKeyVersion);
+    testBranchKeyVersionHappyCase(
+      keyStore := keyStore,
+      branchKeyId := identifier,
+      branchKeyIdActiveVersion := branchKeyVersion,
+      branchKeyIdActiveVersionUtf8Bytes := activeResult.branchKeyMaterials.branchKeyVersion,
+      encryptionContext := encryptionContext
+    );
     VerifyGetKeysFromStorage(identifier, storage);
 
     //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
@@ -93,8 +108,6 @@ module TestGetKeys {
   {
     var ddbClient :- expect DDB.DynamoDBClient();
 
-    print "Running MRK";
-
     var eastKeyStoreConfig := Types.KeyStoreConfig(
       id := None,
       kmsConfiguration := KmsConfigEast,
@@ -413,7 +426,11 @@ module TestGetKeys {
     expect versionResult.error == Types.KeyStoreException(message := ErrorMessages.GET_KEY_ARN_DISAGREEMENT);
   }
 
-  method testBeaconKeyHappyCase(keyStore: Types.IKeyStoreClient, branchKeyId: string)
+  method testBeaconKeyHappyCase(
+    keyStore: Types.IKeyStoreClient,
+    branchKeyId: string,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
+  )
     requires keyStore.ValidState()
     modifies keyStore.Modifies
   {
@@ -421,16 +438,26 @@ module TestGetKeys {
       Types.GetBeaconKeyInput(
         branchKeyIdentifier := branchKeyId
       ));
-    expect isValidBeaconResult?(beaconKeyResult, branchKeyId);
+    expect isValidBeaconResult?(beaconKeyResult, branchKeyId, encryptionContext);
   }
 
-  predicate method isValidBeaconResult?(beaconKeyResult: Types.GetBeaconKeyOutput, branchKeyId: string) {
+  predicate method isValidBeaconResult?(
+    beaconKeyResult: Types.GetBeaconKeyOutput,
+    branchKeyId: string,
+    encryptionContext : Types.EncryptionContext
+  ) {
     && beaconKeyResult.beaconKeyMaterials.beaconKeyIdentifier == branchKeyId
     && beaconKeyResult.beaconKeyMaterials.beaconKey.Some?
     && |beaconKeyResult.beaconKeyMaterials.beaconKey.value| == 32
+    && beaconKeyResult.beaconKeyMaterials.encryptionContext == encryptionContext
   }
 
-  method testActiveBranchKeyHappyCase(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersionUtf8Bytes: seq<uint8>)
+  method testActiveBranchKeyHappyCase(
+    keyStore: Types.IKeyStoreClient,
+    branchKeyId: string,
+    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
+  )
     requires keyStore.ValidState()
     modifies keyStore.Modifies
   {
@@ -438,10 +465,15 @@ module TestGetKeys {
       Types.GetActiveBranchKeyInput(
         branchKeyIdentifier := branchKeyId
       ));
-    expect isValidActiveBranchKeyResult?(branchKeyResult, branchKeyId, Some(branchKeyIdActiveVersionUtf8Bytes));
+    expect isValidActiveBranchKeyResult?(branchKeyResult, branchKeyId, encryptionContext, Some(branchKeyIdActiveVersionUtf8Bytes));
   }
 
-  method testAndGetActiveBranchKeyHappyCase(keyStore: Types.IKeyStoreClient, branchKeyId: string, nameonly branchKeyIdActiveVersionUtf8Bytes: Option<seq<uint8>> := None)
+  method testAndGetActiveBranchKeyHappyCase(
+    keyStore: Types.IKeyStoreClient,
+    branchKeyId: string,
+    nameonly branchKeyIdActiveVersionUtf8Bytes: Option<seq<uint8>> := None,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
+  )
     returns (output: Types.GetActiveBranchKeyOutput)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
@@ -450,18 +482,31 @@ module TestGetKeys {
       Types.GetActiveBranchKeyInput(
         branchKeyIdentifier := branchKeyId
       ));
-    expect isValidActiveBranchKeyResult?(branchKeyResult, branchKeyId, branchKeyIdActiveVersionUtf8Bytes);
+    expect isValidActiveBranchKeyResult?(branchKeyResult, branchKeyId, encryptionContext, branchKeyIdActiveVersionUtf8Bytes);
     return branchKeyResult;
   }
 
-  predicate method isValidActiveBranchKeyResult?(branchKeyResult: Types.GetActiveBranchKeyOutput, branchKeyId: string, branchKeyIdActiveVersionUtf8Bytes: Option<seq<uint8>>) {
+  predicate method isValidActiveBranchKeyResult?(
+    branchKeyResult: Types.GetActiveBranchKeyOutput,
+    branchKeyId: string,
+    encryptionContext : Types.EncryptionContext,
+    branchKeyIdActiveVersionUtf8Bytes: Option<seq<uint8>>
+  )
+  {
     && branchKeyResult.branchKeyMaterials.branchKeyIdentifier == branchKeyId
     && (branchKeyIdActiveVersionUtf8Bytes.None? ||
         branchKeyResult.branchKeyMaterials.branchKeyVersion == branchKeyIdActiveVersionUtf8Bytes.value)
     && |branchKeyResult.branchKeyMaterials.branchKey| == 32
+    && branchKeyResult.branchKeyMaterials.encryptionContext == encryptionContext
   }
 
-  method testBranchKeyVersionHappyCase(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersion: string, branchKeyIdActiveVersionUtf8Bytes: seq<uint8>)
+  method testBranchKeyVersionHappyCase(
+    keyStore: Types.IKeyStoreClient,
+    branchKeyId: string,
+    branchKeyIdActiveVersion: string,
+    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
+  )
     requires keyStore.ValidState()
     modifies keyStore.Modifies
   {
@@ -471,13 +516,21 @@ module TestGetKeys {
         branchKeyVersion := branchKeyIdActiveVersion
       ));
     var testBytes :- expect UTF8.Encode(branchKeyIdActiveVersion);
-    expect isValidBranchKeyVersionResult?(versionResult, branchKeyId, branchKeyIdActiveVersionUtf8Bytes, testBytes);
+    expect isValidBranchKeyVersionResult?(versionResult, branchKeyId, encryptionContext, branchKeyIdActiveVersionUtf8Bytes, testBytes);
   }
 
-  predicate method isValidBranchKeyVersionResult?(versionResult: Types.GetBranchKeyVersionOutput, branchKeyId: string, branchKeyIdActiveVersionUtf8Bytes: seq<uint8>, testBytes: UTF8.ValidUTF8Bytes) {
+  predicate method isValidBranchKeyVersionResult?(
+    versionResult: Types.GetBranchKeyVersionOutput,
+    branchKeyId: string,
+    encryptionContext : Types.EncryptionContext,
+    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>,
+    testBytes: UTF8.ValidUTF8Bytes
+  )
+  {
     && versionResult.branchKeyMaterials.branchKeyIdentifier == branchKeyId
     && versionResult.branchKeyMaterials.branchKeyVersion == branchKeyIdActiveVersionUtf8Bytes == testBytes
     && |versionResult.branchKeyMaterials.branchKey| == 32
+    && versionResult.branchKeyMaterials.encryptionContext == encryptionContext
   }
 
   method VerifyGetKeysFromStorage(
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
index a2d49bcef8..fd1c1044e4 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
@@ -90,6 +90,7 @@ module {:options "/functionSyntax:4" } TestAdminHV2 {
     // Get branch key items from storage
     TestGetKeys.VerifyGetKeys(
       identifier := branchKeyId,
+      encryptionContext := customEC,
       keyStore := keyStore,
       storage := storage
     );
@@ -190,6 +191,7 @@ module {:options "/functionSyntax:4" } TestAdminHV2 {
     // Get branch key items from storage
     TestGetKeys.VerifyGetKeys(
       identifier := createKeyOutput?.Identifier,
+      encryptionContext := customEC,
       keyStore := keyStore,
       storage := storage
     );

From 0bb35b6889aa42d874c33189538785c948ff6d7d Mon Sep 17 00:00:00 2001
From: Rishav karanjit <karanjitrishav4@gmail.com>
Date: Sat, 5 Apr 2025 14:53:00 -0700
Subject: [PATCH 07/13] chore(dafny): refactor HV1 MRK test to use helper
 methods (#1399)

---
 .../AwsCryptographyKeyStore/test/Fixtures.dfy |  57 ++++-
 .../test/TestGetKeys.dfy                      | 199 +++++++++---------
 2 files changed, 153 insertions(+), 103 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
index da87f4a30d..08b8a5b5ec 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
@@ -100,7 +100,19 @@ module Fixtures {
   const KmsMrkConfigAP : Types.KMSConfiguration := Types.KMSConfiguration.kmsMRKeyArn(MrkArnAP)
   const KmsMrkEC : Types.EncryptionContext := map[UTF8.EncodeAscii("abc") := UTF8.EncodeAscii("123")]
   const EastBranchKey : string := "MyEastBranch2"
+  const EastBranchKeyIdActiveVersion : string := "6f22825b-bd56-4434-83e2-2782e2160172"
+  const EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes: seq<uint8> := [
+    54, 102, 50, 50, 56, 50, 53, 98, 45, 98, 100,
+    53, 54, 45, 52, 52, 51, 52, 45, 56, 51, 101, 50,
+    45, 50, 55, 56, 50, 101, 50, 49, 54, 48, 49, 55, 50
+  ]
   const WestBranchKey : string := "MyWestBranch2"
+  const WestBranchKeyIdActiveVersion : string := "094715a4-b98d-4c98-bf50-17422a8938f4"
+  const WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes: seq<uint8> := [
+    48, 57, 52, 55, 49, 53, 97, 52, 45, 98, 57, 56,
+    100, 45, 52, 99, 57, 56, 45, 98, 102, 53, 48, 45,
+    49, 55, 52, 50, 50, 97, 56, 57, 51, 56, 102, 52
+  ]
   const publicKeyArn := "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
 
   // TODO: After ~2024/06/11 launch, add the next two lines to cfn/ESDK-Hierarchy-CI.yaml
@@ -243,14 +255,19 @@ module Fixtures {
     return Success(keyStore);
   }
 
-  method KeyStoreWithNoClient(
+  method KeyStoreWithOptionalClient(
     nameonly kmsId: string := keyArn,
     nameonly physicalName: string := branchKeyStoreName,
-    nameonly logicalName: string := logicalKeyStoreName
+    nameonly logicalName: string := logicalKeyStoreName,
+    nameonly ddbClient?: Option<DDB.Types.IDynamoDBClient> := None,
+    nameonly kmsClient?: Option<KMS.Types.IKMSClient> := None,
+    nameonly srkKey: bool := true,
+    nameonly mrkKey: bool := false
   )
     returns (output: Result<Types.IKeyStoreClient, Types.Error>)
     requires DDB.Types.IsValid_TableName(physicalName)
     requires KMS.Types.IsValid_KeyIdType(kmsId)
+    requires srkKey != mrkKey
     ensures output.Success? ==> output.value.ValidState()
     ensures output.Success?
             ==>
@@ -258,7 +275,19 @@ module Fixtures {
               && fresh(output.value)
               && fresh(output.value.Modifies)
   {
-    var kmsConfig := Types.KMSConfiguration.kmsKeyArn(kmsId);
+    if ddbClient?.Some? {
+      assume {:axiom} fresh(ddbClient?.value) && fresh(ddbClient?.value.Modifies);
+    }
+    if kmsClient?.Some? {
+      assume {:axiom} fresh(kmsClient?.value) && fresh(kmsClient?.value.Modifies);
+    }
+    expect srkKey != mrkKey;
+    var kmsConfig := if srkKey then
+      createSrkKMSConfig(kmsId)
+    else
+      createMrkKMSConfig(kmsId);
+    expect srkKey ==> kmsConfig.kmsKeyArn?;
+    expect mrkKey ==> kmsConfig.kmsMRKeyArn?;
     var keyStoreConfig := Types.KeyStoreConfig(
       id := None,
       kmsConfiguration := kmsConfig,
@@ -266,13 +295,33 @@ module Fixtures {
       storage := Some(
         Types.ddb(
           Types.DynamoDBTable(
-            ddbTableName := physicalName
+            ddbTableName := physicalName,
+            ddbClient := ddbClient?
+          ))),
+      keyManagement := Some(
+        Types.kms(
+          Types.AwsKms(
+            kmsClient := kmsClient?
           )))
     );
     var keyStore :- expect KeyStore.KeyStore(keyStoreConfig);
     return Success(keyStore);
   }
 
+  function method createSrkKMSConfig(kmsId: string) : (output: Types.KMSConfiguration)
+    requires KMS.Types.IsValid_KeyIdType(kmsId)
+    ensures output.kmsKeyArn?
+  {
+    Types.KMSConfiguration.kmsKeyArn(kmsId)
+  }
+
+  function method createMrkKMSConfig(kmsId: string) : (output: Types.KMSConfiguration)
+    requires KMS.Types.IsValid_KeyIdType(kmsId)
+    ensures output.kmsMRKeyArn?
+  {
+    Types.KMSConfiguration.kmsMRKeyArn(kmsId)
+  }
+
   datatype allThree = | allThree (
     active: Types.EncryptedHierarchicalKey,
     beacon: Types.EncryptedHierarchicalKey,
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
index 7860558c52..dda7b7c4aa 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
@@ -108,128 +108,89 @@ module TestGetKeys {
   {
     var ddbClient :- expect DDB.DynamoDBClient();
 
-    var eastKeyStoreConfig := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsConfigEast,
-      logicalKeyStoreName := logicalKeyStoreName,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          )))
+    var westKeyStore :- expect KeyStoreWithOptionalClient(
+      kmsId := MrkArnWest,
+      physicalName := branchKeyStoreName,
+      logicalName := logicalKeyStoreName,
+      ddbClient? := Some(ddbClient)
     );
 
-    var westKeyStoreConfig := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsConfigWest,
-      logicalKeyStoreName := logicalKeyStoreName,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          )))
+    var eastKeyStore :- expect KeyStoreWithOptionalClient(
+      kmsId := MrkArnEast,
+      physicalName := branchKeyStoreName,
+      logicalName := logicalKeyStoreName,
+      ddbClient? := Some(ddbClient)
     );
 
-    var eastMrkKeyStoreConfig := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsMrkConfigEast,
-      logicalKeyStoreName := logicalKeyStoreName,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          )))
+    var westMrkKeyStore :- expect KeyStoreWithOptionalClient(
+      kmsId := MrkArnWest,
+      physicalName := branchKeyStoreName,
+      logicalName := logicalKeyStoreName,
+      ddbClient? := Some(ddbClient),
+      srkKey := false,
+      mrkKey := true
     );
 
-    var westMrkKeyStoreConfig := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsMrkConfigWest,
-      logicalKeyStoreName := logicalKeyStoreName,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          )))
+    var eastMrkKeyStore :- expect KeyStoreWithOptionalClient(
+      kmsId := MrkArnEast,
+      physicalName := branchKeyStoreName,
+      logicalName := logicalKeyStoreName,
+      ddbClient? := Some(ddbClient),
+      srkKey := false,
+      mrkKey := true
     );
 
-    // KmsMrkConfigAP is NOT created
-    var apMrkKeyStoreConfig := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsMrkConfigAP,
-      logicalKeyStoreName := logicalKeyStoreName,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          )))
+    var apMrkKeyStore :- expect KeyStoreWithOptionalClient(
+      kmsId := MrkArnAP,
+      physicalName := branchKeyStoreName,
+      logicalName := logicalKeyStoreName,
+      ddbClient? := Some(ddbClient),
+      srkKey := false,
+      mrkKey := true
     );
 
+    // All four set of keys (branch, beacon and version) should work when the regions match
+    testActiveBranchKeyHappyCase(westKeyStore, WestBranchKey, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testBeaconKeyHappyCase(westKeyStore, WestBranchKey);
+    testBranchKeyVersionHappyCase(westKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
 
-    var westKeyStore :- expect KeyStore.KeyStore(westKeyStoreConfig);
-    var eastKeyStore :- expect KeyStore.KeyStore(eastKeyStoreConfig);
-    var westMrkKeyStore :- expect KeyStore.KeyStore(westMrkKeyStoreConfig);
-    var eastMrkKeyStore :- expect KeyStore.KeyStore(eastMrkKeyStoreConfig);
-    var apMrkKeyStore :- expect KeyStore.KeyStore(apMrkKeyStoreConfig);
+    testActiveBranchKeyHappyCase(eastKeyStore, EastBranchKey, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testBeaconKeyHappyCase(eastKeyStore, EastBranchKey);
+    testBranchKeyVersionHappyCase(eastKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
 
-    // All four should work when the regions match
+    testActiveBranchKeyHappyCase(westMrkKeyStore, WestBranchKey, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testBeaconKeyHappyCase(westMrkKeyStore, WestBranchKey);
+    testBranchKeyVersionHappyCase(westMrkKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
 
-    var activeResult :- expect westKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := WestBranchKey));
-    expect activeResult.branchKeyMaterials.branchKeyIdentifier == WestBranchKey;
-    expect |activeResult.branchKeyMaterials.branchKey| == 32;
-
-    activeResult :- expect eastKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := EastBranchKey));
-    expect activeResult.branchKeyMaterials.branchKeyIdentifier == EastBranchKey;
-    expect |activeResult.branchKeyMaterials.branchKey| == 32;
-
-    activeResult :- expect westMrkKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := WestBranchKey));
-    expect activeResult.branchKeyMaterials.branchKeyIdentifier == WestBranchKey;
-    expect |activeResult.branchKeyMaterials.branchKey| == 32;
-
-    activeResult :- expect eastMrkKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := EastBranchKey));
-    expect activeResult.branchKeyMaterials.branchKeyIdentifier == EastBranchKey;
-    expect |activeResult.branchKeyMaterials.branchKey| == 32;
+    testActiveBranchKeyHappyCase(eastMrkKeyStore, EastBranchKey, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testBeaconKeyHappyCase(eastMrkKeyStore, EastBranchKey);
+    testBranchKeyVersionHappyCase(eastMrkKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
 
     // MRK Configuration should work with the other region
 
-    activeResult :- expect westMrkKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := EastBranchKey));
-    expect activeResult.branchKeyMaterials.branchKeyIdentifier == EastBranchKey;
-    expect |activeResult.branchKeyMaterials.branchKey| == 32;
+    testActiveBranchKeyHappyCase(westMrkKeyStore, EastBranchKey, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testBeaconKeyHappyCase(westMrkKeyStore, EastBranchKey);
+    testBranchKeyVersionHappyCase(westMrkKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
 
-    activeResult :- expect eastMrkKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := WestBranchKey));
-    expect activeResult.branchKeyMaterials.branchKeyIdentifier == WestBranchKey;
-    expect |activeResult.branchKeyMaterials.branchKey| == 32;
+    testActiveBranchKeyHappyCase(eastMrkKeyStore, WestBranchKey, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testBeaconKeyHappyCase(eastMrkKeyStore, WestBranchKey);
+    testBranchKeyVersionHappyCase(eastMrkKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
 
     // Plain Configuration should fail with the other region
 
-    var badResult := westKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := EastBranchKey));
-    expect badResult.Failure?;
-    expect badResult.error == Types.Error.KeyStoreException(message := ErrorMessages.GET_KEY_ARN_DISAGREEMENT);
+    GetActiveKeyWithIncorrectKmsKeyArnHelper(westKeyStore, EastBranchKey);
+    GetBeaconKeyWithIncorrectKmsKeyArnHelper(westKeyStore, EastBranchKey);
+    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(westKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion);
 
-    badResult := eastKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := WestBranchKey));
-    expect badResult.Failure?;
-    expect badResult.error == Types.Error.KeyStoreException(message := ErrorMessages.GET_KEY_ARN_DISAGREEMENT);
+    GetActiveKeyWithIncorrectKmsKeyArnHelper(eastKeyStore, WestBranchKey);
+    GetBeaconKeyWithIncorrectKmsKeyArnHelper(eastKeyStore, WestBranchKey);
+    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(eastKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion);
 
     // apMrkKeyStore should always fail
 
-    badResult := apMrkKeyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(branchKeyIdentifier := WestBranchKey));
-    expect badResult.Failure?;
-    expect badResult.error.ComAmazonawsKms?;
-    expect badResult.error.ComAmazonawsKms.OpaqueWithText?;
-    // it's an opaque error, so I can't test its contents
+    testActiveBranchKeyKMSFailureCase(apMrkKeyStore, WestBranchKey);
+    testBranchKeyVersionKMSFailureCase(apMrkKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion);
+    testBeaconKeyKMSFailureCase(apMrkKeyStore, WestBranchKey);
   }
 
   method {:test} TestKeyWithIncorrectKmsKeyArn() {
@@ -277,7 +238,7 @@ module TestGetKeys {
   method {:test} TestGetKeysWithNoClients() {
     var kmsConfig := Types.KMSConfiguration.kmsKeyArn(keyArn);
 
-    var keyStore :- expect KeyStoreWithNoClient(kmsId:=keyArn, physicalName:=branchKeyStoreName, logicalName := logicalKeyStoreName);
+    var keyStore :- expect KeyStoreWithOptionalClient(kmsId:=keyArn, physicalName:=branchKeyStoreName, logicalName := logicalKeyStoreName);
 
     testActiveBranchKeyHappyCase(keyStore, branchKeyId, branchKeyIdActiveVersionUtf8Bytes);
     testActiveBranchKeyHappyCase(keyStore, hv2BranchKeyId, hv2BranchKeyIdActiveVersionUtf8Bytes);
@@ -533,6 +494,46 @@ module TestGetKeys {
     && versionResult.branchKeyMaterials.encryptionContext == encryptionContext
   }
 
+  method testActiveBranchKeyKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string)
+    requires keyStore.ValidState()
+    modifies keyStore.Modifies
+  {
+    var branchKeyResult := keyStore.GetActiveBranchKey(
+      Types.GetActiveBranchKeyInput(
+        branchKeyIdentifier := branchKeyId
+      ));
+    expect branchKeyResult.Failure?;
+    expect branchKeyResult.error.ComAmazonawsKms?;
+    expect branchKeyResult.error.ComAmazonawsKms.OpaqueWithText?;
+  }
+
+  method testBranchKeyVersionKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersion: string)
+    requires keyStore.ValidState()
+    modifies keyStore.Modifies
+  {
+    var versionResult := keyStore.GetBranchKeyVersion(
+      Types.GetBranchKeyVersionInput(
+        branchKeyIdentifier := branchKeyId,
+        branchKeyVersion := branchKeyIdActiveVersion
+      ));
+    expect versionResult.Failure?;
+    expect versionResult.error.ComAmazonawsKms?;
+    expect versionResult.error.ComAmazonawsKms.OpaqueWithText?;
+  }
+
+  method testBeaconKeyKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string)
+    requires keyStore.ValidState()
+    modifies keyStore.Modifies
+  {
+    var beaconKeyResult := keyStore.GetBeaconKey(
+      Types.GetBeaconKeyInput(
+        branchKeyIdentifier := branchKeyId
+      ));
+    expect beaconKeyResult.Failure?;
+    expect beaconKeyResult.error.ComAmazonawsKms?;
+    expect beaconKeyResult.error.ComAmazonawsKms.OpaqueWithText?;
+  }
+
   method VerifyGetKeysFromStorage(
     identifier : string,
     storage : Types.IKeyStorageInterface

From 4c4e06d8085b0fd3f6f77c1c512b8e89bb186d7d Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Sun, 6 Apr 2025 04:40:37 -0700
Subject: [PATCH 08/13] ensure EC is tested

---
 .../AwsCryptographyKeyStore/test/Fixtures.dfy |  15 +-
 .../test/TestGetKeys.dfy                      | 178 +++++++++++++++---
 .../test/TestAdminCreateKeys.dfy              |   2 +
 3 files changed, 158 insertions(+), 37 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
index 08b8a5b5ec..2659358f4b 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
@@ -61,6 +61,7 @@ module Fixtures {
 
   const branchKeyStoreName := "KeyStoreDdbTable"
   const logicalKeyStoreName := branchKeyStoreName
+  // hierarchy-version-1 branch key
   const branchKeyId := "3f43a9af-08c5-4317-b694-3d3e883dcaef"
   const branchKeyIdActiveVersion := "a4905627-4b7f-4272-a847-f50dae245737"
   // This is branchKeyIdActiveVersion above, as utf8bytes
@@ -71,14 +72,16 @@ module Fixtures {
     56, 52, 55, 45, 102, 53, 48, 100, 97, 101,
     50, 52, 53, 55, 51, 55
   ]
-
   const branchKeyIdWithEC := "4bb57643-07c1-419e-92ad-0df0df149d7c"
-  const hv2BranchKeyId := "test-hv2-branch-key-badaa332-29f2-4c72-8ad7-071eb48499c3"
-  const hv2BranchKeyVersion := "347fdc7d-e93f-4166-97c2-5f5e0053d335"
+  // hierarchy-version-2 branch key
+  const hv2BranchKeyId := "4a0c7b92-3703-4209-8961-24b07ab6562b"
+  const hv2BranchKeyVersion := "a0496b5c-e048-42bc-8b75-68a004851803"
+  // This is hv2BranchKeyVersion above, as utf8bytes
+  // https://cyberchef.infosec.amazon.dev/#recipe=Encode_text('UTF-8%20(65001)')To_Decimal('Comma',false)&input=YTA0OTZiNWMtZTA0OC00MmJjLThiNzUtNjhhMDA0ODUxODAz&oenc=65001
   const hv2BranchKeyIdActiveVersionUtf8Bytes: seq<uint8> := [
-    51, 52, 55, 102, 100, 99, 55, 100, 45, 101, 57,
-    51, 102, 45, 52, 49, 54, 54, 45, 57, 55, 99, 50,
-    45, 53, 102, 53, 101, 48, 48, 53, 51, 100, 51, 51, 53
+    97, 48, 52, 57, 54, 98, 53, 99, 45, 101, 48, 52, 
+    56, 45, 52, 50, 98, 99, 45, 56, 98, 55, 53, 45, 
+    54, 56, 97, 48, 48, 52, 56, 53, 49, 56, 48, 51
   ]
   // THESE ARE TESTING RESOURCES DO NOT USE IN A PRODUCTION ENVIRONMENT
   const keyArn := "arn:aws:kms:us-west-2:370957321024:key/9d989aa2-2f9c-438c-a745-cc57d3ad0126"
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
index a012fce49e..aa8c2924d3 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
@@ -150,47 +150,163 @@ module TestGetKeys {
     );
 
     // All four set of keys (branch, beacon and version) should work when the regions match
-    testActiveBranchKeyHappyCase(westKeyStore, WestBranchKey, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
-    testBeaconKeyHappyCase(westKeyStore, WestBranchKey);
-    testBranchKeyVersionHappyCase(westKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
-
-    testActiveBranchKeyHappyCase(eastKeyStore, EastBranchKey, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
-    testBeaconKeyHappyCase(eastKeyStore, EastBranchKey);
-    testBranchKeyVersionHappyCase(eastKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testActiveBranchKeyHappyCase(
+      keyStore := westKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
+    testBeaconKeyHappyCase(
+      keyStore := westKeyStore,
+      branchKeyId := WestBranchKey,
+      encryptionContext := KmsMrkEC
+    );
+    testBranchKeyVersionHappyCase(
+      keyStore := westKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion,
+      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
 
-    testActiveBranchKeyHappyCase(westMrkKeyStore, WestBranchKey, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
-    testBeaconKeyHappyCase(westMrkKeyStore, WestBranchKey);
-    testBranchKeyVersionHappyCase(westMrkKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testActiveBranchKeyHappyCase(
+      keyStore := eastKeyStore,
+      branchKeyId := EastBranchKey,
+      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
+    testBeaconKeyHappyCase(
+      keyStore := eastKeyStore,
+      branchKeyId := EastBranchKey,
+      encryptionContext := KmsMrkEC
+    );
+    testBranchKeyVersionHappyCase(
+      keyStore := eastKeyStore,
+      branchKeyId := EastBranchKey,
+      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion,
+      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
 
-    testActiveBranchKeyHappyCase(eastMrkKeyStore, EastBranchKey, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
-    testBeaconKeyHappyCase(eastMrkKeyStore, EastBranchKey);
-    testBranchKeyVersionHappyCase(eastMrkKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testActiveBranchKeyHappyCase(
+      keyStore := westMrkKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
+    testBeaconKeyHappyCase(
+      keyStore := westMrkKeyStore,
+      branchKeyId := WestBranchKey,
+      encryptionContext := KmsMrkEC
+    );
+    testBranchKeyVersionHappyCase(
+      keyStore := westMrkKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion,
+      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
 
-    // MRK Configuration should work with the other region
+    testActiveBranchKeyHappyCase(
+      keyStore := eastMrkKeyStore,
+      branchKeyId := EastBranchKey,
+      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
+    testBeaconKeyHappyCase(
+      keyStore := eastMrkKeyStore,
+      branchKeyId := EastBranchKey,
+      encryptionContext := KmsMrkEC
+    );
+    testBranchKeyVersionHappyCase(
+      keyStore := eastMrkKeyStore,
+      branchKeyId := EastBranchKey,
+      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion,
+      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
 
-    testActiveBranchKeyHappyCase(westMrkKeyStore, EastBranchKey, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
-    testBeaconKeyHappyCase(westMrkKeyStore, EastBranchKey);
-    testBranchKeyVersionHappyCase(westMrkKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion, EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testActiveBranchKeyHappyCase(
+      keyStore := westMrkKeyStore,
+      branchKeyId := EastBranchKey,
+      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
+    testBeaconKeyHappyCase(
+      keyStore := westMrkKeyStore,
+      branchKeyId := EastBranchKey,
+      encryptionContext := KmsMrkEC
+    );
+    testBranchKeyVersionHappyCase(
+      keyStore := westMrkKeyStore,
+      branchKeyId := EastBranchKey,
+      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion,
+      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
 
-    testActiveBranchKeyHappyCase(eastMrkKeyStore, WestBranchKey, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
-    testBeaconKeyHappyCase(eastMrkKeyStore, WestBranchKey);
-    testBranchKeyVersionHappyCase(eastMrkKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion, WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes);
+    testActiveBranchKeyHappyCase(
+      keyStore := eastMrkKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
+    testBeaconKeyHappyCase(
+      keyStore := eastMrkKeyStore,
+      branchKeyId := WestBranchKey,
+      encryptionContext := KmsMrkEC
+    );
+    testBranchKeyVersionHappyCase(
+      keyStore := eastMrkKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion,
+      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
+      encryptionContext := KmsMrkEC
+    );
 
     // Plain Configuration should fail with the other region
+    GetActiveKeyWithIncorrectKmsKeyArnHelper(
+      keyStore := westKeyStore,
+      branchKeyId := EastBranchKey
+    );
+    GetBeaconKeyWithIncorrectKmsKeyArnHelper(
+      keyStore := westKeyStore,
+      branchKeyId := EastBranchKey
+    );
+    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(
+      keyStore := westKeyStore,
+      branchKeyId := EastBranchKey,
+      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion
+    );
 
-    GetActiveKeyWithIncorrectKmsKeyArnHelper(westKeyStore, EastBranchKey);
-    GetBeaconKeyWithIncorrectKmsKeyArnHelper(westKeyStore, EastBranchKey);
-    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(westKeyStore, EastBranchKey, EastBranchKeyIdActiveVersion);
-
-    GetActiveKeyWithIncorrectKmsKeyArnHelper(eastKeyStore, WestBranchKey);
-    GetBeaconKeyWithIncorrectKmsKeyArnHelper(eastKeyStore, WestBranchKey);
-    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(eastKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion);
+    GetActiveKeyWithIncorrectKmsKeyArnHelper(
+      keyStore := eastKeyStore,
+      branchKeyId := WestBranchKey
+    );
+    GetBeaconKeyWithIncorrectKmsKeyArnHelper(
+      keyStore := eastKeyStore,
+      branchKeyId := WestBranchKey
+    );
+    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(
+      keyStore := eastKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion
+    );
 
     // apMrkKeyStore should always fail
-
-    testActiveBranchKeyKMSFailureCase(apMrkKeyStore, WestBranchKey);
-    testBranchKeyVersionKMSFailureCase(apMrkKeyStore, WestBranchKey, WestBranchKeyIdActiveVersion);
-    testBeaconKeyKMSFailureCase(apMrkKeyStore, WestBranchKey);
+    testActiveBranchKeyKMSFailureCase(
+      keyStore := apMrkKeyStore,
+      branchKeyId := WestBranchKey
+    );
+    testBranchKeyVersionKMSFailureCase(
+      keyStore := apMrkKeyStore,
+      branchKeyId := WestBranchKey,
+      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion
+    );
+    testBeaconKeyKMSFailureCase(
+      keyStore := apMrkKeyStore,
+      branchKeyId := WestBranchKey
+    );
   }
 
   method {:test} TestKeyWithIncorrectKmsKeyArn() {
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy
index be966b9035..763e4e9e5c 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy
@@ -41,6 +41,8 @@ module {:options "/functionSyntax:4" } TestAdminCreateKeys {
     );
     var identifier? :- expect underTest.CreateKey(input);
     var identifier := identifier?.Identifier;
+    expect identifier?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v1,
+      "KeyStoreAdmin should create branch key with `hierarchy-version-1` when no `HierarchyVersion` provided";
 
     // Get branch key items from storage
     TestGetKeys.VerifyGetKeys(

From 26a75a33cfdd26322599b31a3e03578d110219dd Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Sun, 6 Apr 2025 04:46:51 -0700
Subject: [PATCH 09/13] nit

---
 .../dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy           | 1 +
 1 file changed, 1 insertion(+)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
index aa8c2924d3..a1448e9dfc 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
@@ -226,6 +226,7 @@ module TestGetKeys {
       encryptionContext := KmsMrkEC
     );
 
+    // MRK Configuration should work with the other region
     testActiveBranchKeyHappyCase(
       keyStore := westMrkKeyStore,
       branchKeyId := EastBranchKey,

From f6bae98ed77c948272f8837ca93b32e593cf89ae Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Sun, 6 Apr 2025 04:56:07 -0700
Subject: [PATCH 10/13] remove duplicate

---
 .../AwsCryptographyKeyStore/test/Fixtures.dfy |  4 +-
 .../test/TestGetKeys.dfy                      | 40 -------------------
 2 files changed, 2 insertions(+), 42 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
index 2659358f4b..a2392b7eb7 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
@@ -79,8 +79,8 @@ module Fixtures {
   // This is hv2BranchKeyVersion above, as utf8bytes
   // https://cyberchef.infosec.amazon.dev/#recipe=Encode_text('UTF-8%20(65001)')To_Decimal('Comma',false)&input=YTA0OTZiNWMtZTA0OC00MmJjLThiNzUtNjhhMDA0ODUxODAz&oenc=65001
   const hv2BranchKeyIdActiveVersionUtf8Bytes: seq<uint8> := [
-    97, 48, 52, 57, 54, 98, 53, 99, 45, 101, 48, 52, 
-    56, 45, 52, 50, 98, 99, 45, 56, 98, 55, 53, 45, 
+    97, 48, 52, 57, 54, 98, 53, 99, 45, 101, 48, 52,
+    56, 45, 52, 50, 98, 99, 45, 56, 98, 55, 53, 45,
     54, 56, 97, 48, 48, 52, 56, 53, 49, 56, 48, 51
   ]
   // THESE ARE TESTING RESOURCES DO NOT USE IN A PRODUCTION ENVIRONMENT
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
index a1448e9dfc..3e45126ef5 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
@@ -651,46 +651,6 @@ module TestGetKeys {
     expect beaconKeyResult.error.ComAmazonawsKms.OpaqueWithText?;
   }
 
-  method testActiveBranchKeyKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string)
-    requires keyStore.ValidState()
-    modifies keyStore.Modifies
-  {
-    var branchKeyResult := keyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(
-        branchKeyIdentifier := branchKeyId
-      ));
-    expect branchKeyResult.Failure?;
-    expect branchKeyResult.error.ComAmazonawsKms?;
-    expect branchKeyResult.error.ComAmazonawsKms.OpaqueWithText?;
-  }
-
-  method testBranchKeyVersionKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersion: string)
-    requires keyStore.ValidState()
-    modifies keyStore.Modifies
-  {
-    var versionResult := keyStore.GetBranchKeyVersion(
-      Types.GetBranchKeyVersionInput(
-        branchKeyIdentifier := branchKeyId,
-        branchKeyVersion := branchKeyIdActiveVersion
-      ));
-    expect versionResult.Failure?;
-    expect versionResult.error.ComAmazonawsKms?;
-    expect versionResult.error.ComAmazonawsKms.OpaqueWithText?;
-  }
-
-  method testBeaconKeyKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string)
-    requires keyStore.ValidState()
-    modifies keyStore.Modifies
-  {
-    var beaconKeyResult := keyStore.GetBeaconKey(
-      Types.GetBeaconKeyInput(
-        branchKeyIdentifier := branchKeyId
-      ));
-    expect beaconKeyResult.Failure?;
-    expect beaconKeyResult.error.ComAmazonawsKms?;
-    expect beaconKeyResult.error.ComAmazonawsKms.OpaqueWithText?;
-  }
-
   method VerifyGetKeysFromStorage(
     identifier : string,
     storage : Types.IKeyStorageInterface

From 536077f69b14f663e866bbe59bd6a70671a4596e Mon Sep 17 00:00:00 2001
From: Darwin Chowdary <39110935+imabhichow@users.noreply.github.com>
Date: Mon, 7 Apr 2025 09:47:49 -0700
Subject: [PATCH 11/13] isolate assertions

---
 .../dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
index 500a6d4995..1877cc4a25 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
@@ -25,7 +25,7 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
 
   const happyCaseId := "test-happy-case-create-key-hv-1"
 
-  method {:test} TestCreateMRKForHV1()
+  method {:test} {:isolate_assertions} TestCreateMRKForHV1()
   {
     var ddbClient :- expect DDB.DynamoDBClient();
 

From bfbdad1433543c6a93a65b668211a50120f7b7df Mon Sep 17 00:00:00 2001
From: Tony Knapp <5892063+texastony@users.noreply.github.com>
Date: Tue, 8 Apr 2025 08:06:07 -0700
Subject: [PATCH 12/13] test(dafny): refactor BKS & BKSA tests (#1406)

---
 .../test/BranchKeyValidators.dfy              | 256 ++++++++
 .../AwsCryptographyKeyStore/test/Fixtures.dfy |  41 +-
 .../test/TestConfig.dfy                       |  22 +
 .../test/TestCreateKeys.dfy                   | 133 ++---
 .../test/TestGetKeys.dfy                      | 553 +++---------------
 .../test/TestAdminCreateKeys.dfy              | 165 +++++-
 .../test/TestAdminHV1Only.dfy                 | 123 ++++
 .../test/TestAdminHV2.dfy                     | 326 -----------
 8 files changed, 675 insertions(+), 944 deletions(-)
 create mode 100644 AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy
 create mode 100644 AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
 delete mode 100644 AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy
new file mode 100644
index 0000000000..6ad5163cd4
--- /dev/null
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy
@@ -0,0 +1,256 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+include "../src/Index.dfy"
+
+// methods to validate a Get* Branch Key result from the Branch Key Store client
+module {:options "/functionSyntax:4" } BranchKeyValidators {
+  import opened Wrappers
+  import opened StandardLibrary.UInt
+  import UTF8
+  import UUID
+  import Types = AwsCryptographyKeyStoreTypes
+  import Structure
+
+  method VerifyGetKeys(
+    identifier : string,
+    keyStore : Types.IKeyStoreClient,
+    storage : Types.IKeyStorageInterface,
+    nameonly versionUtf8Bytes?: Option<seq<uint8>> := None,
+    nameonly encryptionContext : Types.EncryptionContext := map[],
+    nameonly hierarchyVersion : Types.HierarchyVersion := Types.HierarchyVersion.v1
+  )
+    requires
+      keyStore.ValidState() && storage.ValidState()
+    modifies
+      keyStore.Modifies, storage.Modifies
+    ensures
+      keyStore.ValidState() && storage.ValidState()
+
+  {
+    var _ := testBeaconKeyHappyCase(
+      keyStore := keyStore,
+      branchKeyId := identifier,
+      encryptionContext := encryptionContext
+    );
+
+    var activeResult := testActiveBranchKeyHappyCase(
+      keyStore := keyStore,
+      branchKeyId := identifier,
+      encryptionContext := encryptionContext
+    );
+    var versionString :- expect UTF8.Decode(activeResult.branchKeyVersion);
+
+    var _ := testBranchKeyVersionHappyCase(
+      keyStore := keyStore,
+      branchKeyId := identifier,
+      branchKeyIdActiveVersion := versionString,
+      branchKeyIdActiveVersionUtf8Bytes := activeResult.branchKeyVersion,
+      encryptionContext := encryptionContext
+    );
+    VerifyGetKeysFromStorage(identifier, storage, hierarchyVersion:=hierarchyVersion);
+
+    //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
+    //= type=test
+    //# This guid MUST be [version 4 UUID](https://www.ietf.org/rfc/rfc4122.txt)
+    var versionByteUUID :- expect UUID.ToByteArray(versionString);
+    var versionRoundTrip :- expect UUID.FromByteArray(versionByteUUID);
+    expect versionRoundTrip == versionString;
+  }
+
+
+  method VerifyGetKeysFromStorage(
+    identifier : string,
+    storage : Types.IKeyStorageInterface,
+    nameonly hierarchyVersion : Types.HierarchyVersion := Types.HierarchyVersion.v1
+  )
+    requires storage.ValidState()
+    modifies storage.Modifies
+    ensures storage.ValidState()
+  {
+    var active :- expect storage.GetEncryptedActiveBranchKey(
+      Types.GetEncryptedActiveBranchKeyInput(
+        Identifier := identifier
+      )
+    );
+    expect active.Item.Type.ActiveHierarchicalSymmetricVersion?;
+
+    var beacon :- expect storage.GetEncryptedBeaconKey(
+      Types.GetEncryptedBeaconKeyInput(
+        Identifier := identifier
+      )
+    );
+    expect beacon.Item.Type.ActiveHierarchicalSymmetricBeacon?;
+
+    var version :- expect storage.GetEncryptedBranchKeyVersion(
+      Types.GetEncryptedBranchKeyVersionInput(
+        Identifier := identifier,
+        Version := active.Item.Type.ActiveHierarchicalSymmetricVersion.Version
+      )
+    );
+    expect version.Item.Type.HierarchicalSymmetricVersion?;
+
+    var hvMatches? := match hierarchyVersion {
+      case v1
+        =>
+        && branchKeyContextSaysHV1(active.Item.EncryptionContext)
+        && branchKeyContextSaysHV1(beacon.Item.EncryptionContext)
+        && branchKeyContextSaysHV1(version.Item.EncryptionContext)
+      case v2
+        =>
+        && branchKeyContextSaysHV2(active.Item.EncryptionContext)
+        && branchKeyContextSaysHV2(beacon.Item.EncryptionContext)
+        && branchKeyContextSaysHV2(version.Item.EncryptionContext)
+    };
+    if (!hvMatches?) {
+      print "HV did not match expectation, did bkc have hv?: ", Structure.HIERARCHY_VERSION in active.Item.EncryptionContext;
+      if (Structure.HIERARCHY_VERSION in active.Item.EncryptionContext) {
+        print " Actual HV: ", active.Item.EncryptionContext[Structure.HIERARCHY_VERSION];
+      }
+    }
+    expect hvMatches?, "Hierarchy Version did not match expectation.";
+
+    //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
+    //= type=test
+    //# This timestamp MUST be in ISO 8601 format in UTC, to microsecond precision (e.g. “YYYY-MM-DDTHH:mm:ss.ssssssZ“)
+    expect ISO8601?(active.Item.CreateTime);
+    expect ISO8601?(beacon.Item.CreateTime);
+    expect ISO8601?(version.Item.CreateTime);
+  }
+
+  predicate branchKeyContextSaysHV1(bkc: Types.EncryptionContextString)
+  {
+    Structure.HIERARCHY_VERSION in bkc && bkc[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_VALUE_1
+  }
+
+  predicate branchKeyContextSaysHV2(bkc: Types.EncryptionContextString)
+  {
+    Structure.HIERARCHY_VERSION in bkc && bkc[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_VALUE_2
+  }
+
+  method testActiveBranchKeyHappyCase(
+    keyStore: Types.IKeyStoreClient,
+    branchKeyId: string,
+    nameonly versionUtf8Bytes?: Option<seq<uint8>> := None,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
+  ) returns (output: Types.BranchKeyMaterials)
+    requires keyStore.ValidState()
+    modifies keyStore.Modifies
+    ensures keyStore.ValidState()
+  {
+    var activeResult :- expect keyStore.GetActiveBranchKey(
+      Types.GetActiveBranchKeyInput(
+        branchKeyIdentifier := branchKeyId
+      ));
+    expect isValidActiveBranchKeyResult?(activeResult, branchKeyId, encryptionContext, versionUtf8Bytes?);
+    return activeResult.branchKeyMaterials;
+  }
+
+  method testBranchKeyVersionHappyCase(
+    keyStore: Types.IKeyStoreClient,
+    branchKeyId: string,
+    branchKeyIdActiveVersion: string,
+    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
+  ) returns (output: Types.BranchKeyMaterials)
+    requires keyStore.ValidState()
+    modifies keyStore.Modifies
+    ensures keyStore.ValidState()
+  {
+    var versionResult :- expect keyStore.GetBranchKeyVersion(
+      Types.GetBranchKeyVersionInput(
+        branchKeyIdentifier := branchKeyId,
+        branchKeyVersion := branchKeyIdActiveVersion
+      ));
+    expect isValidBranchKeyVersionResult?(versionResult, branchKeyId, encryptionContext, branchKeyIdActiveVersionUtf8Bytes);
+    return versionResult.branchKeyMaterials;
+  }
+
+  method testBeaconKeyHappyCase(
+    keyStore: Types.IKeyStoreClient,
+    branchKeyId: string,
+    nameonly encryptionContext : Types.EncryptionContext := map[]
+  ) returns (output: Types.BeaconKeyMaterials)
+    requires keyStore.ValidState()
+    modifies keyStore.Modifies
+    ensures keyStore.ValidState()
+  {
+    var beaconKeyResult :- expect keyStore.GetBeaconKey(
+      Types.GetBeaconKeyInput(
+        branchKeyIdentifier := branchKeyId
+      ));
+    expect beaconKeyResult.beaconKeyMaterials.beaconKeyIdentifier == branchKeyId, "Wrong Branch Key ID.";
+    expect beaconKeyResult.beaconKeyMaterials.beaconKey.Some?, "Beacon Key Materials MUST be present.";
+    expect |beaconKeyResult.beaconKeyMaterials.beaconKey.value| == 32, "Lenght of Beacon Key Materials MUST be 32 bytes.";
+    if (beaconKeyResult.beaconKeyMaterials.encryptionContext != encryptionContext) {
+      print "Beacon Key's encryption context is incorrect. Expected: ", encryptionContext, " but got: ", beaconKeyResult.beaconKeyMaterials.encryptionContext;
+    }
+    expect beaconKeyResult.beaconKeyMaterials.encryptionContext == encryptionContext, "Beacon Key's encryption context is incorrect.";
+    expect isValidBeaconResult?(beaconKeyResult, branchKeyId, encryptionContext);
+    return beaconKeyResult.beaconKeyMaterials;
+  }
+
+  predicate isValidBeaconResult?(
+    beaconKeyResult: Types.GetBeaconKeyOutput,
+    branchKeyId: string,
+    encryptionContext : Types.EncryptionContext
+  ) {
+    && beaconKeyResult.beaconKeyMaterials.beaconKeyIdentifier == branchKeyId
+    && beaconKeyResult.beaconKeyMaterials.beaconKey.Some?
+    && |beaconKeyResult.beaconKeyMaterials.beaconKey.value| == 32
+    && beaconKeyResult.beaconKeyMaterials.encryptionContext == encryptionContext
+  }
+
+  predicate isValidActiveBranchKeyResult?(
+    branchKeyResult: Types.GetActiveBranchKeyOutput,
+    branchKeyId: string,
+    encryptionContext : Types.EncryptionContext,
+    branchKeyIdActiveVersionUtf8Bytes: Option<seq<uint8>>
+  )
+  {
+    && branchKeyResult.branchKeyMaterials.branchKeyIdentifier == branchKeyId
+    && (branchKeyIdActiveVersionUtf8Bytes.None? ||
+        branchKeyResult.branchKeyMaterials.branchKeyVersion == branchKeyIdActiveVersionUtf8Bytes.value)
+    && |branchKeyResult.branchKeyMaterials.branchKey| == 32
+    && branchKeyResult.branchKeyMaterials.encryptionContext == encryptionContext
+  }
+
+  predicate isValidBranchKeyVersionResult?(
+    versionResult: Types.GetBranchKeyVersionOutput,
+    branchKeyId: string,
+    encryptionContext : Types.EncryptionContext,
+    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>
+  )
+  {
+    && versionResult.branchKeyMaterials.branchKeyIdentifier == branchKeyId
+    && versionResult.branchKeyMaterials.branchKeyVersion == branchKeyIdActiveVersionUtf8Bytes
+    && |versionResult.branchKeyMaterials.branchKey| == 32
+    && versionResult.branchKeyMaterials.encryptionContext == encryptionContext
+  }
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
+  //= type=test
+  //# This timestamp MUST be in ISO 8601 format in UTC, to microsecond precision (e.g. “YYYY-MM-DDTHH:mm:ss.ssssssZ“)
+  lemma ISO8601Test()
+  {
+    assert ISO8601?("2024-08-06T17:23:25.000874Z");
+  }
+
+  //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
+  //= type=test
+  //# This timestamp MUST be in ISO 8601 format in UTC, to microsecond precision (e.g. “YYYY-MM-DDTHH:mm:ss.ssssssZ“)
+  predicate ISO8601?(
+    CreateTime: string
+  )
+  {
+    // “YYYY-MM-DDTHH:mm:ss.ssssssZ“
+    && |CreateTime| == 27
+    && CreateTime[4] == '-'
+    && CreateTime[7] == '-'
+    && CreateTime[10] == 'T'
+    && CreateTime[13] == ':'
+    && CreateTime[16] == ':'
+    && CreateTime[19] == '.'
+    && CreateTime[26] == 'Z'
+  }
+}
\ No newline at end of file
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
index a2392b7eb7..04300a7131 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/Fixtures.dfy
@@ -102,6 +102,7 @@ module Fixtures {
   const KmsSrkConfigWest : Types.KMSConfiguration := Types.KMSConfiguration.kmsKeyArn(MrkArnWest)
   const KmsMrkConfigAP : Types.KMSConfiguration := Types.KMSConfiguration.kmsMRKeyArn(MrkArnAP)
   const KmsMrkEC : Types.EncryptionContext := map[UTF8.EncodeAscii("abc") := UTF8.EncodeAscii("123")]
+  const RobbieEC : Types.EncryptionContext := map[UTF8.EncodeAscii("Robbie") := UTF8.EncodeAscii("is a dog.")]
   const EastBranchKey : string := "MyEastBranch2"
   const EastBranchKeyIdActiveVersion : string := "6f22825b-bd56-4434-83e2-2782e2160172"
   const EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes: seq<uint8> := [
@@ -258,19 +259,14 @@ module Fixtures {
     return Success(keyStore);
   }
 
-  method KeyStoreWithOptionalClient(
-    nameonly kmsId: string := keyArn,
+  method KeyStoreFromKMSConfig(
+    nameonly kmsConfig: Types.KMSConfiguration,
     nameonly physicalName: string := branchKeyStoreName,
     nameonly logicalName: string := logicalKeyStoreName,
-    nameonly ddbClient?: Option<DDB.Types.IDynamoDBClient> := None,
-    nameonly kmsClient?: Option<KMS.Types.IKMSClient> := None,
-    nameonly srkKey: bool := true,
-    nameonly mrkKey: bool := false
+    nameonly ddbClient?: Option<DDB.Types.IDynamoDBClient> := None
   )
     returns (output: Result<Types.IKeyStoreClient, Types.Error>)
     requires DDB.Types.IsValid_TableName(physicalName)
-    requires KMS.Types.IsValid_KeyIdType(kmsId)
-    requires srkKey != mrkKey
     ensures output.Success? ==> output.value.ValidState()
     ensures output.Success?
             ==>
@@ -281,16 +277,6 @@ module Fixtures {
     if ddbClient?.Some? {
       assume {:axiom} fresh(ddbClient?.value) && fresh(ddbClient?.value.Modifies);
     }
-    if kmsClient?.Some? {
-      assume {:axiom} fresh(kmsClient?.value) && fresh(kmsClient?.value.Modifies);
-    }
-    expect srkKey != mrkKey;
-    var kmsConfig := if srkKey then
-      createSrkKMSConfig(kmsId)
-    else
-      createMrkKMSConfig(kmsId);
-    expect srkKey ==> kmsConfig.kmsKeyArn?;
-    expect mrkKey ==> kmsConfig.kmsMRKeyArn?;
     var keyStoreConfig := Types.KeyStoreConfig(
       id := None,
       kmsConfiguration := kmsConfig,
@@ -300,31 +286,12 @@ module Fixtures {
           Types.DynamoDBTable(
             ddbTableName := physicalName,
             ddbClient := ddbClient?
-          ))),
-      keyManagement := Some(
-        Types.kms(
-          Types.AwsKms(
-            kmsClient := kmsClient?
           )))
     );
     var keyStore :- expect KeyStore.KeyStore(keyStoreConfig);
     return Success(keyStore);
   }
 
-  function method createSrkKMSConfig(kmsId: string) : (output: Types.KMSConfiguration)
-    requires KMS.Types.IsValid_KeyIdType(kmsId)
-    ensures output.kmsKeyArn?
-  {
-    Types.KMSConfiguration.kmsKeyArn(kmsId)
-  }
-
-  function method createMrkKMSConfig(kmsId: string) : (output: Types.KMSConfiguration)
-    requires KMS.Types.IsValid_KeyIdType(kmsId)
-    ensures output.kmsMRKeyArn?
-  {
-    Types.KMSConfiguration.kmsMRKeyArn(kmsId)
-  }
-
   datatype allThree = | allThree (
     active: Types.EncryptedHierarchicalKey,
     beacon: Types.EncryptedHierarchicalKey,
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestConfig.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestConfig.dfy
index 5f96008f85..bd70caddfd 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestConfig.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestConfig.dfy
@@ -3,7 +3,10 @@
 
 include "../src/Index.dfy"
 include "Fixtures.dfy"
+include "BranchKeyValidators.dfy"
 
+// This really tests the logic in BKS/Index.dfy
+// and KmsArn.dfy
 module TestConfig {
   import Types = AwsCryptographyKeyStoreTypes
   import ComAmazonawsKmsTypes
@@ -14,6 +17,7 @@ module TestConfig {
   import opened Fixtures
   import UUID
   import ErrorMessages = KeyStoreErrorMessages
+  import BranchKeyValidators
 
   method {:test} TestInvalidKmsKeyArnConfig() {
     var kmsClient :- expect KMS.KMSClient();
@@ -188,4 +192,22 @@ module TestConfig {
     keyStore := KeyStore.KeyStore(keyStoreConfig);
     expect keyStore.Success?;
   }
+
+  // Actually tests constructing a Key Store with no Clients
+  method {:test} TestGetKeysWithNoClients() {
+    var kmsConfig := Types.KMSConfiguration.kmsKeyArn(keyArn);
+    var keyStoreConfig := Types.KeyStoreConfig(
+      id := None,
+      kmsConfiguration := kmsConfig,
+      logicalKeyStoreName := logicalKeyStoreName,
+      grantTokens := None,
+      ddbTableName := Some(branchKeyStoreName),
+      ddbClient := None,
+      kmsClient := None
+    );
+    var keyStore :- expect KeyStore.KeyStore(keyStoreConfig);
+    var _ := BranchKeyValidators.testBeaconKeyHappyCase(keyStore, branchKeyId);
+    var _ := BranchKeyValidators.testActiveBranchKeyHappyCase(keyStore, branchKeyId, versionUtf8Bytes?:=Some(branchKeyIdActiveVersionUtf8Bytes));
+    var _ := BranchKeyValidators.testBranchKeyVersionHappyCase(keyStore, branchKeyId, branchKeyIdActiveVersion, branchKeyIdActiveVersionUtf8Bytes);
+  }
 }
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
index 1877cc4a25..06a0c5d5fd 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestCreateKeys.dfy
@@ -4,7 +4,8 @@
 include "../src/Index.dfy"
 include "Fixtures.dfy"
 include "CleanupItems.dfy"
-include "TestGetKeys.dfy"
+include "BranchKeyValidators.dfy"
+
 
 module {:options "/functionSyntax:4" } TestCreateKeys {
   import Types = AwsCryptographyKeyStoreTypes
@@ -21,108 +22,58 @@ module {:options "/functionSyntax:4" } TestCreateKeys {
   import KmsArn
   import UUID
   import AwsArnParsing
-  import TestGetKeys
+  import BranchKeyValidators
 
   const happyCaseId := "test-happy-case-create-key-hv-1"
-
-  method {:test} {:isolate_assertions} TestCreateMRKForHV1()
+  method {:test} TestCreateMRKForHV1()
   {
     var ddbClient :- expect DDB.DynamoDBClient();
-
-    var keyStoreConfigEast := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsConfigEast,
-      logicalKeyStoreName := logicalKeyStoreName,
-      grantTokens := None,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          )))
+    var storage :- expect DefaultStorage(ddbClient? := Some(ddbClient));
+    var keyStoreEast :- expect KeyStoreFromKMSConfig(
+      kmsConfig := KmsConfigEast,
+      ddbClient? := Some(ddbClient)
     );
-
-    var keyStoreConfigWest := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := KmsConfigWest,
-      logicalKeyStoreName := logicalKeyStoreName,
-      grantTokens := None,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          )))
-    );
-
-    // Create key with Custom EC & Branch Key Identifier
     var uuid :- expect UUID.GenerateUUID();
-    var branchKeyIdWest := happyCaseId + "-" + "west" + "-" + uuid;
-    // print branchKeyIdWest;
-    var branchKeyIdEast := happyCaseId + "-" + "east" + "-" + uuid;
-    // print branchKeyIdEast;
-
-    var keyStoreEast :- expect KeyStore.KeyStore(keyStoreConfigEast);
-    var keyStoreWest :- expect KeyStore.KeyStore(keyStoreConfigWest);
-
-    var branchKeyIdWest? :- expect keyStoreWest.CreateKey(Types.CreateKeyInput(
-                                                            branchKeyIdentifier := Some(branchKeyIdWest),
-                                                            encryptionContext := Some(KmsMrkEC)
-                                                          ));
-
-    var branchKeyIdEast? :- expect keyStoreEast.CreateKey(Types.CreateKeyInput(
-                                                            branchKeyIdentifier := Some(branchKeyIdEast),
-                                                            encryptionContext := Some(KmsMrkEC)
-                                                          ));
-
-    expect branchKeyIdEast?.branchKeyIdentifier == branchKeyIdEast;
-    expect branchKeyIdWest?.branchKeyIdentifier == branchKeyIdWest;
-
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyIdWest, ddbClient:=ddbClient);
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyIdEast, ddbClient:=ddbClient);
+    var bkid := happyCaseId + "-east-mrk-" + uuid;
+    var bk :- expect keyStoreEast.CreateKey(
+      Types.CreateKeyInput(
+        branchKeyIdentifier := Some(bkid),
+        encryptionContext := Some(KmsMrkEC)
+      ));
+    expect bk.branchKeyIdentifier == bkid;
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStoreEast, storage,
+                                      encryptionContext := KmsMrkEC);
+    var keyStoreWest :- expect KeyStoreFromKMSConfig(
+      // Passing in KmsMrkConfigWest causes the KMS Client to be created for us-west-2
+      // EastBranchKey's KMS-ARN is replicated to us-west-2
+      kmsConfig := KmsMrkConfigWest,
+      ddbClient? := Some(ddbClient)
+    );
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStoreWest, storage,
+                                      encryptionContext := KmsMrkEC);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=bkid, ddbClient:=ddbClient);
   }
 
-  method {:test} TestCreateBranchAndBeaconKeys()
+  method {:test} TestCreateSRKForHV1()
   {
-    var kmsClient :- expect KMS.KMSClient();
     var ddbClient :- expect DDB.DynamoDBClient();
+    var storage :- expect DefaultStorage(ddbClient? := Some(ddbClient));
     var kmsConfig := Types.KMSConfiguration.kmsKeyArn(keyArn);
-
-    var keyStoreConfig := Types.KeyStoreConfig(
-      id := None,
-      kmsConfiguration := kmsConfig,
-      logicalKeyStoreName := logicalKeyStoreName,
-      storage := Some(
-        Types.ddb(
-          Types.DynamoDBTable(
-            ddbTableName := branchKeyStoreName,
-            ddbClient := Some(ddbClient)
-          ))),
-      keyManagement := Some(
-        Types.kms(
-          Types.AwsKms(
-            kmsClient := Some(kmsClient)
-          )))
+    var keyStore :- expect KeyStoreFromKMSConfig(
+      kmsConfig := kmsConfig,
+      ddbClient? := Some(ddbClient)
     );
-
-    var keyStore :- expect KeyStore.KeyStore(keyStoreConfig);
-
-    var branchKeyId :- expect keyStore.CreateKey(Types.CreateKeyInput(
-                                                   branchKeyIdentifier := None,
-                                                   encryptionContext := None
-                                                 ));
-
-    // Get branch key items from storage
-    TestGetKeys.VerifyGetKeys(
-      identifier := branchKeyId.branchKeyIdentifier,
-      keyStore := keyStore,
-      storage := keyStore.config.storage
-    );
-
-    // Since this process uses a read DDB table,
-    // the number of records will forever increase.
-    // To avoid this, remove the items.
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyId.branchKeyIdentifier, ddbClient:=ddbClient);
+    var uuid :- expect UUID.GenerateUUID();
+    var bkid := happyCaseId + "-srk-" + uuid;
+    var bk :- expect keyStore.CreateKey(
+      Types.CreateKeyInput(
+        branchKeyIdentifier := Some(bkid),
+        encryptionContext := Some(KmsMrkEC)
+      ));
+    expect bk.branchKeyIdentifier == bkid;
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStore, storage,
+                                      encryptionContext := KmsMrkEC);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=bkid, ddbClient:=ddbClient);
   }
 
   method {:test} TestCreateOptions()
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
index 3e45126ef5..169ac991f5 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/TestGetKeys.dfy
@@ -5,10 +5,9 @@ include "../src/Index.dfy"
 include "Fixtures.dfy"
 include "CleanupItems.dfy"
 include "../src/ErrorMessages.dfy"
+include "BranchKeyValidators.dfy"
 
-// TODO-HV2-M1: add more HV2 tests.
-// TODO-HV2-M1: Maybe rename this module to TestGetHv1Keys and create another TestGetHv2Keys
-module TestGetKeys {
+module {:options "/functionSyntax:4" } TestGetKeys {
   import Types = AwsCryptographyKeyStoreTypes
   import ComAmazonawsKmsTypes
   import KMS = Com.Amazonaws.Kms
@@ -21,355 +20,124 @@ module TestGetKeys {
   import UTF8
   import ErrorMessages = KeyStoreErrorMessages
   import UUID
+  import BranchKeyValidators
 
   const incorrectLogicalName := "MySuperAwesomeTableName"
 
-  method VerifyGetKeys(
-    identifier : string,
-    keyStore : Types.IKeyStoreClient,
-    storage : Types.IKeyStorageInterface,
-    nameonly encryptionContext : Types.EncryptionContext := map[]
-  )
-    requires
-      keyStore.ValidState() && storage.ValidState()
-    modifies
-      keyStore.Modifies, storage.Modifies
-    ensures
-      keyStore.ValidState() && storage.ValidState()
-
+  method {:test} TestGetKeysHappyCase()
   {
-    testBeaconKeyHappyCase(
-      keyStore := keyStore,
-      branchKeyId := identifier,
-      encryptionContext := encryptionContext
-    );
-
-    var activeResult := testAndGetActiveBranchKeyHappyCase(
-      keyStore := keyStore,
-      branchKeyId := identifier,
-      encryptionContext := encryptionContext
-    );
-    var branchKeyVersion :- expect UTF8.Decode(activeResult.branchKeyMaterials.branchKeyVersion);
-
-    testBranchKeyVersionHappyCase(
-      keyStore := keyStore,
-      branchKeyId := identifier,
-      branchKeyIdActiveVersion := branchKeyVersion,
-      branchKeyIdActiveVersionUtf8Bytes := activeResult.branchKeyMaterials.branchKeyVersion,
-      encryptionContext := encryptionContext
-    );
-    VerifyGetKeysFromStorage(identifier, storage);
-
-    //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
-    //= type=test
-    //# This guid MUST be [version 4 UUID](https://www.ietf.org/rfc/rfc4122.txt)
-    var versionString :- expect UTF8.Decode(activeResult.branchKeyMaterials.branchKeyVersion);
-    var versionByteUUID :- expect UUID.ToByteArray(versionString);
-    var versionRoundTrip :- expect UUID.FromByteArray(versionByteUUID);
-    expect versionRoundTrip == versionString;
+    var ddbClient :- expect DDB.DynamoDBClient();
+    var keyStore :- expect DefaultKeyStore(kmsId := keyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName, ddbClient? := Some(ddbClient));
+    var storage :- expect DefaultStorage(ddbClient? := Some(ddbClient));
+    BranchKeyValidators.VerifyGetKeys(branchKeyId, keyStore, storage, versionUtf8Bytes?:=Some(branchKeyIdActiveVersionUtf8Bytes));
   }
 
-  method {:test} TestGetKeysHappyCase()
+  method {:test} TestGetKeysHV2HappyCase()
   {
-    var kmsClient :- expect KMS.KMSClient();
     var ddbClient :- expect DDB.DynamoDBClient();
-    var kmsConfig := Types.KMSConfiguration.kmsKeyArn(keyArn);
-
-    var keyStore :- expect DefaultKeyStore(kmsId := keyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName, ddbClient? := Some(ddbClient), kmsClient? := Some(kmsClient));
-
-    testBeaconKeyHappyCase(keyStore, branchKeyId);
-    testBeaconKeyHappyCase(keyStore, hv2BranchKeyId);
-
-    testActiveBranchKeyHappyCase(keyStore, branchKeyId, branchKeyIdActiveVersionUtf8Bytes);
-    testActiveBranchKeyHappyCase(keyStore, hv2BranchKeyId, hv2BranchKeyIdActiveVersionUtf8Bytes);
-
-    testBranchKeyVersionHappyCase(keyStore, branchKeyId, branchKeyIdActiveVersion, branchKeyIdActiveVersionUtf8Bytes);
-    testBranchKeyVersionHappyCase(keyStore, hv2BranchKeyId, hv2BranchKeyVersion, hv2BranchKeyIdActiveVersionUtf8Bytes);
+    var keyStore :- expect DefaultKeyStore(kmsId := keyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName, ddbClient? := Some(ddbClient));
+    var storage :- expect DefaultStorage(ddbClient? := Some(ddbClient));
+    BranchKeyValidators.VerifyGetKeys(hv2BranchKeyId, keyStore, storage,
+                                      versionUtf8Bytes?:=Some(hv2BranchKeyIdActiveVersionUtf8Bytes),
+                                      hierarchyVersion := Types.HierarchyVersion.v2);
   }
 
-  // This is a static test case to Get Branch Keys created with Mrk Keys
-  method {:test} TestGetActiveMrkKey()
+  method {:test} TestGetMRKeySameRegionHappyCase()
   {
-    VerifyGetActiveMrkKey(
-      KmsConfigEast := KmsConfigEast,
-      KmsConfigWest := KmsConfigWest,
-      KmsMrkConfigEast := KmsMrkConfigEast,
-      KmsMrkConfigWest := KmsMrkConfigWest
+    var ddbClient :- expect DDB.DynamoDBClient();
+    var keyStore :- expect KeyStoreFromKMSConfig(
+      kmsConfig := KmsMrkConfigWest,
+      ddbClient? := Some(ddbClient)
     );
+    var storage :- expect DefaultStorage(ddbClient? := Some(ddbClient));
+    BranchKeyValidators.VerifyGetKeys(WestBranchKey, keyStore, storage,
+                                      versionUtf8Bytes?:=Some(WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes),
+                                      encryptionContext := KmsMrkEC);
   }
 
-  // TODO-HV2-M1: Add MRK test for hv2
-  method {:isolate_assertions} VerifyGetActiveMrkKey(
-    KmsConfigEast : Types.KMSConfiguration,
-    KmsConfigWest : Types.KMSConfiguration,
-    KmsMrkConfigEast : Types.KMSConfiguration,
-    KmsMrkConfigWest : Types.KMSConfiguration
-  )
+  method {:test} TestGetMRKeyReplicatedRegionHappyCase()
   {
     var ddbClient :- expect DDB.DynamoDBClient();
-
-    var westKeyStore :- expect KeyStoreWithOptionalClient(
-      kmsId := MrkArnWest,
-      physicalName := branchKeyStoreName,
-      logicalName := logicalKeyStoreName,
+    var storage :- expect DefaultStorage(ddbClient? := Some(ddbClient));
+    var keyStore :- expect KeyStoreFromKMSConfig(
+      // Passing in KmsMrkConfigWest causes the KMS Client to be created for us-west-2
+      // EastBranchKey's KMS-ARN is replicated to us-west-2
+      kmsConfig := KmsMrkConfigWest,
       ddbClient? := Some(ddbClient)
     );
+    BranchKeyValidators.VerifyGetKeys(EastBranchKey, keyStore, storage,
+                                      versionUtf8Bytes?:=Some(EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes),
+                                      encryptionContext := KmsMrkEC);
+  }
 
-    var eastKeyStore :- expect KeyStoreWithOptionalClient(
-      kmsId := MrkArnEast,
-      physicalName := branchKeyStoreName,
-      logicalName := logicalKeyStoreName,
+  method {:test} TestGetMRKeyMrkNotReplicatedFails()
+  {
+    var ddbClient :- expect DDB.DynamoDBClient();
+    var keyStore :- expect KeyStoreFromKMSConfig(
+      // Passing in KmsMrkConfigAP causes the KMS Client to be created for ap-south-2
+      // EastBranchKey's KMS-ARN is NOT replicated to ap-south-2
+      kmsConfig := KmsMrkConfigAP,
+      // However, we MUST pass in a DDB Client, as that needs to be created in us-west-2,
+      // or the error will be a DDB Error.
       ddbClient? := Some(ddbClient)
     );
+    testActiveBranchKeyKMSFailureCase(keyStore, EastBranchKey);
+    testBeaconKeyKMSFailureCase(keyStore, EastBranchKey);
+    testBranchKeyVersionKMSFailureCase(keyStore, EastBranchKey, EastBranchKeyIdActiveVersion);
+  }
 
-    var westMrkKeyStore :- expect KeyStoreWithOptionalClient(
-      kmsId := MrkArnWest,
-      physicalName := branchKeyStoreName,
-      logicalName := logicalKeyStoreName,
-      ddbClient? := Some(ddbClient),
-      srkKey := false,
-      mrkKey := true
-    );
-
-    var eastMrkKeyStore :- expect KeyStoreWithOptionalClient(
-      kmsId := MrkArnEast,
-      physicalName := branchKeyStoreName,
-      logicalName := logicalKeyStoreName,
-      ddbClient? := Some(ddbClient),
-      srkKey := false,
-      mrkKey := true
-    );
-
-    var apMrkKeyStore :- expect KeyStoreWithOptionalClient(
-      kmsId := MrkArnAP,
-      physicalName := branchKeyStoreName,
-      logicalName := logicalKeyStoreName,
-      ddbClient? := Some(ddbClient),
-      srkKey := false,
-      mrkKey := true
-    );
-
-    // All four set of keys (branch, beacon and version) should work when the regions match
-    testActiveBranchKeyHappyCase(
-      keyStore := westKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-    testBeaconKeyHappyCase(
-      keyStore := westKeyStore,
-      branchKeyId := WestBranchKey,
-      encryptionContext := KmsMrkEC
-    );
-    testBranchKeyVersionHappyCase(
-      keyStore := westKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion,
-      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-
-    testActiveBranchKeyHappyCase(
-      keyStore := eastKeyStore,
-      branchKeyId := EastBranchKey,
-      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-    testBeaconKeyHappyCase(
-      keyStore := eastKeyStore,
-      branchKeyId := EastBranchKey,
-      encryptionContext := KmsMrkEC
-    );
-    testBranchKeyVersionHappyCase(
-      keyStore := eastKeyStore,
-      branchKeyId := EastBranchKey,
-      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion,
-      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-
-    testActiveBranchKeyHappyCase(
-      keyStore := westMrkKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-    testBeaconKeyHappyCase(
-      keyStore := westMrkKeyStore,
-      branchKeyId := WestBranchKey,
-      encryptionContext := KmsMrkEC
-    );
-    testBranchKeyVersionHappyCase(
-      keyStore := westMrkKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion,
-      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-
-    testActiveBranchKeyHappyCase(
-      keyStore := eastMrkKeyStore,
-      branchKeyId := EastBranchKey,
-      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-    testBeaconKeyHappyCase(
-      keyStore := eastMrkKeyStore,
-      branchKeyId := EastBranchKey,
-      encryptionContext := KmsMrkEC
-    );
-    testBranchKeyVersionHappyCase(
-      keyStore := eastMrkKeyStore,
-      branchKeyId := EastBranchKey,
-      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion,
-      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-
-    // MRK Configuration should work with the other region
-    testActiveBranchKeyHappyCase(
-      keyStore := westMrkKeyStore,
-      branchKeyId := EastBranchKey,
-      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-    testBeaconKeyHappyCase(
-      keyStore := westMrkKeyStore,
-      branchKeyId := EastBranchKey,
-      encryptionContext := KmsMrkEC
-    );
-    testBranchKeyVersionHappyCase(
-      keyStore := westMrkKeyStore,
-      branchKeyId := EastBranchKey,
-      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion,
-      branchKeyIdActiveVersionUtf8Bytes := EastBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-
-    testActiveBranchKeyHappyCase(
-      keyStore := eastMrkKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-    testBeaconKeyHappyCase(
-      keyStore := eastMrkKeyStore,
-      branchKeyId := WestBranchKey,
-      encryptionContext := KmsMrkEC
-    );
-    testBranchKeyVersionHappyCase(
-      keyStore := eastMrkKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion,
-      branchKeyIdActiveVersionUtf8Bytes := WestBranchKeyBranchKeyIdActiveVersionUtf8Bytes,
-      encryptionContext := KmsMrkEC
-    );
-
-    // Plain Configuration should fail with the other region
-    GetActiveKeyWithIncorrectKmsKeyArnHelper(
-      keyStore := westKeyStore,
-      branchKeyId := EastBranchKey
-    );
-    GetBeaconKeyWithIncorrectKmsKeyArnHelper(
-      keyStore := westKeyStore,
-      branchKeyId := EastBranchKey
-    );
-    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(
-      keyStore := westKeyStore,
-      branchKeyId := EastBranchKey,
-      branchKeyIdActiveVersion := EastBranchKeyIdActiveVersion
-    );
-
-    GetActiveKeyWithIncorrectKmsKeyArnHelper(
-      keyStore := eastKeyStore,
-      branchKeyId := WestBranchKey
-    );
-    GetBeaconKeyWithIncorrectKmsKeyArnHelper(
-      keyStore := eastKeyStore,
-      branchKeyId := WestBranchKey
-    );
-    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(
-      keyStore := eastKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion
-    );
-
-    // apMrkKeyStore should always fail
-    testActiveBranchKeyKMSFailureCase(
-      keyStore := apMrkKeyStore,
-      branchKeyId := WestBranchKey
-    );
-    testBranchKeyVersionKMSFailureCase(
-      keyStore := apMrkKeyStore,
-      branchKeyId := WestBranchKey,
-      branchKeyIdActiveVersion := WestBranchKeyIdActiveVersion
-    );
-    testBeaconKeyKMSFailureCase(
-      keyStore := apMrkKeyStore,
-      branchKeyId := WestBranchKey
+  method {:test} TestGetSRKeyButKmsArnOutOfRegionFails()
+  {
+    var keyStore :- expect KeyStoreFromKMSConfig(
+      kmsConfig := KmsConfigWest
     );
+    GetActiveKeyWithIncorrectKmsKeyArnHelper(keyStore, EastBranchKey);
+    GetBeaconKeyWithIncorrectKmsKeyArnHelper(keyStore, EastBranchKey);
+    GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(keyStore, EastBranchKey, EastBranchKeyIdActiveVersion);
   }
 
   method {:test} TestKeyWithIncorrectKmsKeyArn() {
-    var kmsClient :- expect KMS.KMSClient();
-    var ddbClient :- expect DDB.DynamoDBClient();
-
-    var keyStore :- expect DefaultKeyStore(kmsId := postalHornKeyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName, ddbClient? := Some(ddbClient), kmsClient? := Some(kmsClient));
-
+    var keyStore :- expect DefaultKeyStore(kmsId := postalHornKeyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName);
     GetActiveKeyWithIncorrectKmsKeyArnHelper(keyStore, branchKeyId);
-    GetActiveKeyWithIncorrectKmsKeyArnHelper(keyStore, hv2BranchKeyId);
-
     GetBeaconKeyWithIncorrectKmsKeyArnHelper(keyStore, branchKeyId);
-    GetBeaconKeyWithIncorrectKmsKeyArnHelper(keyStore, hv2BranchKeyId);
-
     GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(keyStore, branchKeyId, branchKeyIdActiveVersion);
+  }
+
+  method {:test} TestKeyWithIncorrectKmsKeyArnHV2() {
+    var keyStore :- expect DefaultKeyStore(kmsId := postalHornKeyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName);
+    GetActiveKeyWithIncorrectKmsKeyArnHelper(keyStore, hv2BranchKeyId);
+    GetBeaconKeyWithIncorrectKmsKeyArnHelper(keyStore, hv2BranchKeyId);
     GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(keyStore, hv2BranchKeyId, hv2BranchKeyVersion);
   }
 
   method {:test} TestGetActiveKeyWrongLogicalKeyStoreName() {
-    var kmsClient :- expect KMS.KMSClient();
-    var ddbClient :- expect DDB.DynamoDBClient();
-
-    var keyStore :- expect DefaultKeyStore(kmsId:=keyArn, physicalName := branchKeyStoreName, logicalName := incorrectLogicalName, ddbClient? := Some(ddbClient), kmsClient? := Some(kmsClient));
+    var keyStore :- expect DefaultKeyStore(kmsId:=keyArn, physicalName := branchKeyStoreName, logicalName := incorrectLogicalName);
     GetActiveKeyWrongLogicalKeyStoreName(keyStore, branchKeyId, Types.HierarchyVersion.v1);
-    GetActiveKeyWrongLogicalKeyStoreName(keyStore, hv2BranchKeyId, Types.HierarchyVersion.v2);
-
     GetBeaconKeyWrongLogicalKeyStoreName(keyStore, branchKeyId, Types.HierarchyVersion.v1);
-    GetBeaconKeyWrongLogicalKeyStoreName(keyStore, hv2BranchKeyId, Types.HierarchyVersion.v2);
-
     GetVersionKeyWrongLogicalKeyStoreName(keyStore, branchKeyId, branchKeyIdActiveVersion, Types.HierarchyVersion.v1);
+  }
+
+
+  method {:test} TestGetActiveKeyWrongLogicalKeyStoreNameHV2() {
+    var keyStore :- expect DefaultKeyStore(kmsId:=keyArn, physicalName := branchKeyStoreName, logicalName := incorrectLogicalName);
+    GetActiveKeyWrongLogicalKeyStoreName(keyStore, hv2BranchKeyId, Types.HierarchyVersion.v2);
+    GetBeaconKeyWrongLogicalKeyStoreName(keyStore, hv2BranchKeyId, Types.HierarchyVersion.v2);
     GetVersionKeyWrongLogicalKeyStoreName(keyStore, hv2BranchKeyId, hv2BranchKeyVersion, Types.HierarchyVersion.v2);
   }
 
+  // This test does not consider HV b/c it is really testing the storage layer
   method {:test} TestGetKeyDoesNotExistFails()
   {
-    var kmsClient :- expect KMS.KMSClient();
-    var ddbClient :- expect DDB.DynamoDBClient();
-
-    var keyStore :- expect DefaultKeyStore(kmsId := keyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName, ddbClient? := Some(ddbClient), kmsClient? := Some(kmsClient));
+    var keyStore :- expect DefaultKeyStore(kmsId := keyArn, physicalName := branchKeyStoreName, logicalName := logicalKeyStoreName);
     GetActiveKeyDoesNotExistFailsHelper(keyStore, "Robbie");
     GetBeaconKeyDoesNotExistFailsHelper(keyStore, "Robbie");
     GetBranchKeyVersionDoesNotExistFailsHelper(keyStore, "Robbie", branchKeyIdActiveVersion);
   }
 
-  method {:test} TestGetKeysWithNoClients() {
-    var kmsConfig := Types.KMSConfiguration.kmsKeyArn(keyArn);
-
-    var keyStore :- expect KeyStoreWithOptionalClient(kmsId:=keyArn, physicalName:=branchKeyStoreName, logicalName := logicalKeyStoreName);
-
-    testActiveBranchKeyHappyCase(keyStore, branchKeyId, branchKeyIdActiveVersionUtf8Bytes);
-    testActiveBranchKeyHappyCase(keyStore, hv2BranchKeyId, hv2BranchKeyIdActiveVersionUtf8Bytes);
-
-    testBeaconKeyHappyCase(keyStore, branchKeyId);
-    testBeaconKeyHappyCase(keyStore, hv2BranchKeyId);
-
-    testBranchKeyVersionHappyCase(keyStore, branchKeyId, branchKeyIdActiveVersion, branchKeyIdActiveVersionUtf8Bytes);
-    testBranchKeyVersionHappyCase(keyStore, hv2BranchKeyId, hv2BranchKeyVersion, hv2BranchKeyIdActiveVersionUtf8Bytes);
-  }
-
   method GetActiveKeyWrongLogicalKeyStoreName(keyStore: Types.IKeyStoreClient, branchKeyId: string, hv: Types.HierarchyVersion)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var activeResult := keyStore.GetActiveBranchKey(
       Types.GetActiveBranchKeyInput(
@@ -390,6 +158,7 @@ module TestGetKeys {
   method GetBeaconKeyWrongLogicalKeyStoreName(keyStore: Types.IKeyStoreClient, branchKeyId: string, hv: Types.HierarchyVersion)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var beaconKeyResult := keyStore.GetBeaconKey(
       Types.GetBeaconKeyInput(
@@ -410,6 +179,7 @@ module TestGetKeys {
   method GetVersionKeyWrongLogicalKeyStoreName(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersion: string, hv: Types.HierarchyVersion)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var versionResult := keyStore.GetBranchKeyVersion(
       Types.GetBranchKeyVersionInput(
@@ -431,6 +201,7 @@ module TestGetKeys {
   method GetActiveKeyDoesNotExistFailsHelper(keyStore: Types.IKeyStoreClient, branchKeyId: string)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var activeResult := keyStore.GetActiveBranchKey(
       Types.GetActiveBranchKeyInput(
@@ -443,6 +214,7 @@ module TestGetKeys {
   method GetBeaconKeyDoesNotExistFailsHelper(keyStore: Types.IKeyStoreClient, branchKeyId: string)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var beaconKeyResult := keyStore.GetBeaconKey(
       Types.GetBeaconKeyInput(
@@ -455,6 +227,7 @@ module TestGetKeys {
   method GetBranchKeyVersionDoesNotExistFailsHelper(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersion: string)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var versionResult := keyStore.GetBranchKeyVersion(
       Types.GetBranchKeyVersionInput(
@@ -474,12 +247,16 @@ module TestGetKeys {
         branchKeyIdentifier := branchKeyId
       ));
     expect activeResult.Failure?;
+    if (activeResult.error != Types.KeyStoreException(message := ErrorMessages.GET_KEY_ARN_DISAGREEMENT)) {
+      print "Expected GET_KEY_ARN_DISAGREEMENT but got: ", activeResult.error;
+    }
     expect activeResult.error == Types.KeyStoreException(message := ErrorMessages.GET_KEY_ARN_DISAGREEMENT);
   }
 
   method GetBeaconKeyWithIncorrectKmsKeyArnHelper(keyStore: Types.IKeyStoreClient, branchKeyId: string)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var beaconKeyResult := keyStore.GetBeaconKey(
       Types.GetBeaconKeyInput(
@@ -493,6 +270,7 @@ module TestGetKeys {
   method GetBranchKeyVersionWithIncorrectKmsKeyArnHelper(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersion: string)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var versionResult := keyStore.GetBranchKeyVersion(
       Types.GetBranchKeyVersionInput(
@@ -504,122 +282,19 @@ module TestGetKeys {
     expect versionResult.error == Types.KeyStoreException(message := ErrorMessages.GET_KEY_ARN_DISAGREEMENT);
   }
 
-  method testBeaconKeyHappyCase(
-    keyStore: Types.IKeyStoreClient,
-    branchKeyId: string,
-    nameonly encryptionContext : Types.EncryptionContext := map[]
-  )
-    requires keyStore.ValidState()
-    modifies keyStore.Modifies
-  {
-    var beaconKeyResult :- expect keyStore.GetBeaconKey(
-      Types.GetBeaconKeyInput(
-        branchKeyIdentifier := branchKeyId
-      ));
-    expect isValidBeaconResult?(beaconKeyResult, branchKeyId, encryptionContext);
-  }
-
-  predicate method isValidBeaconResult?(
-    beaconKeyResult: Types.GetBeaconKeyOutput,
-    branchKeyId: string,
-    encryptionContext : Types.EncryptionContext
-  ) {
-    && beaconKeyResult.beaconKeyMaterials.beaconKeyIdentifier == branchKeyId
-    && beaconKeyResult.beaconKeyMaterials.beaconKey.Some?
-    && |beaconKeyResult.beaconKeyMaterials.beaconKey.value| == 32
-    && beaconKeyResult.beaconKeyMaterials.encryptionContext == encryptionContext
-  }
-
-  method testActiveBranchKeyHappyCase(
-    keyStore: Types.IKeyStoreClient,
-    branchKeyId: string,
-    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>,
-    nameonly encryptionContext : Types.EncryptionContext := map[]
-  )
-    requires keyStore.ValidState()
-    modifies keyStore.Modifies
-  {
-    var branchKeyResult :- expect keyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(
-        branchKeyIdentifier := branchKeyId
-      ));
-    expect isValidActiveBranchKeyResult?(branchKeyResult, branchKeyId, encryptionContext, Some(branchKeyIdActiveVersionUtf8Bytes));
-  }
-
-  method testAndGetActiveBranchKeyHappyCase(
-    keyStore: Types.IKeyStoreClient,
-    branchKeyId: string,
-    nameonly branchKeyIdActiveVersionUtf8Bytes: Option<seq<uint8>> := None,
-    nameonly encryptionContext : Types.EncryptionContext := map[]
-  )
-    returns (output: Types.GetActiveBranchKeyOutput)
-    requires keyStore.ValidState()
-    modifies keyStore.Modifies
-  {
-    var branchKeyResult :- expect keyStore.GetActiveBranchKey(
-      Types.GetActiveBranchKeyInput(
-        branchKeyIdentifier := branchKeyId
-      ));
-    expect isValidActiveBranchKeyResult?(branchKeyResult, branchKeyId, encryptionContext, branchKeyIdActiveVersionUtf8Bytes);
-    return branchKeyResult;
-  }
-
-  predicate method isValidActiveBranchKeyResult?(
-    branchKeyResult: Types.GetActiveBranchKeyOutput,
-    branchKeyId: string,
-    encryptionContext : Types.EncryptionContext,
-    branchKeyIdActiveVersionUtf8Bytes: Option<seq<uint8>>
-  )
-  {
-    && branchKeyResult.branchKeyMaterials.branchKeyIdentifier == branchKeyId
-    && (branchKeyIdActiveVersionUtf8Bytes.None? ||
-        branchKeyResult.branchKeyMaterials.branchKeyVersion == branchKeyIdActiveVersionUtf8Bytes.value)
-    && |branchKeyResult.branchKeyMaterials.branchKey| == 32
-    && branchKeyResult.branchKeyMaterials.encryptionContext == encryptionContext
-  }
-
-  method testBranchKeyVersionHappyCase(
-    keyStore: Types.IKeyStoreClient,
-    branchKeyId: string,
-    branchKeyIdActiveVersion: string,
-    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>,
-    nameonly encryptionContext : Types.EncryptionContext := map[]
-  )
-    requires keyStore.ValidState()
-    modifies keyStore.Modifies
-  {
-    var versionResult :- expect keyStore.GetBranchKeyVersion(
-      Types.GetBranchKeyVersionInput(
-        branchKeyIdentifier := branchKeyId,
-        branchKeyVersion := branchKeyIdActiveVersion
-      ));
-    var testBytes :- expect UTF8.Encode(branchKeyIdActiveVersion);
-    expect isValidBranchKeyVersionResult?(versionResult, branchKeyId, encryptionContext, branchKeyIdActiveVersionUtf8Bytes, testBytes);
-  }
-
-  predicate method isValidBranchKeyVersionResult?(
-    versionResult: Types.GetBranchKeyVersionOutput,
-    branchKeyId: string,
-    encryptionContext : Types.EncryptionContext,
-    branchKeyIdActiveVersionUtf8Bytes: seq<uint8>,
-    testBytes: UTF8.ValidUTF8Bytes
-  )
-  {
-    && versionResult.branchKeyMaterials.branchKeyIdentifier == branchKeyId
-    && versionResult.branchKeyMaterials.branchKeyVersion == branchKeyIdActiveVersionUtf8Bytes == testBytes
-    && |versionResult.branchKeyMaterials.branchKey| == 32
-    && versionResult.branchKeyMaterials.encryptionContext == encryptionContext
-  }
-
   method testActiveBranchKeyKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var branchKeyResult := keyStore.GetActiveBranchKey(
       Types.GetActiveBranchKeyInput(
         branchKeyIdentifier := branchKeyId
       ));
     expect branchKeyResult.Failure?;
+    if (!branchKeyResult.error.ComAmazonawsKms?) {
+      print "Expected KMS Error, but got: ", branchKeyResult.error;
+    }
     expect branchKeyResult.error.ComAmazonawsKms?;
     expect branchKeyResult.error.ComAmazonawsKms.OpaqueWithText?;
   }
@@ -627,6 +302,7 @@ module TestGetKeys {
   method testBranchKeyVersionKMSFailureCase(keyStore: Types.IKeyStoreClient, branchKeyId: string, branchKeyIdActiveVersion: string)
     requires keyStore.ValidState()
     modifies keyStore.Modifies
+    ensures keyStore.ValidState()
   {
     var versionResult := keyStore.GetBranchKeyVersion(
       Types.GetBranchKeyVersionInput(
@@ -650,61 +326,4 @@ module TestGetKeys {
     expect beaconKeyResult.error.ComAmazonawsKms?;
     expect beaconKeyResult.error.ComAmazonawsKms.OpaqueWithText?;
   }
-
-  method VerifyGetKeysFromStorage(
-    identifier : string,
-    storage : Types.IKeyStorageInterface
-  )
-    requires storage.ValidState()
-    modifies storage.Modifies
-    ensures storage.ValidState()
-  {
-    var encryptedActiveFromStorage :- expect storage.GetEncryptedActiveBranchKey(
-      Types.GetEncryptedActiveBranchKeyInput(
-        Identifier := identifier
-      )
-    );
-
-    var encryptedBeaconFromStorage :- expect storage.GetEncryptedBeaconKey(
-      Types.GetEncryptedBeaconKeyInput(
-        Identifier := identifier
-      )
-    );
-
-    expect encryptedActiveFromStorage.Item.Type.ActiveHierarchicalSymmetricVersion?;
-
-    var encryptedVersionFromStorage :- expect storage.GetEncryptedBranchKeyVersion(
-      Types.GetEncryptedBranchKeyVersionInput(
-        Identifier := identifier,
-        Version := encryptedActiveFromStorage.Item.Type.ActiveHierarchicalSymmetricVersion.Version
-      )
-    );
-
-    //= aws-encryption-sdk-specification/framework/branch-key-store.md#branch-key-and-beacon-key-creation
-    //= type=test
-    //# This timestamp MUST be in ISO 8601 format in UTC, to microsecond precision (e.g. “YYYY-MM-DDTHH:mm:ss.ssssssZ“)
-    expect ISO8601?(encryptedActiveFromStorage.Item.CreateTime);
-    expect ISO8601?(encryptedBeaconFromStorage.Item.CreateTime);
-    expect ISO8601?(encryptedVersionFromStorage.Item.CreateTime);
-  }
-
-  lemma ISO8601Test()
-  {
-    assert ISO8601?("2024-08-06T17:23:25.000874Z");
-  }
-
-  predicate method ISO8601?(
-    CreateTime: string
-  )
-  {
-    // “YYYY-MM-DDTHH:mm:ss.ssssssZ“
-    && |CreateTime| == 27
-    && CreateTime[4] == '-'
-    && CreateTime[7] == '-'
-    && CreateTime[10] == 'T'
-    && CreateTime[13] == ':'
-    && CreateTime[16] == ':'
-    && CreateTime[19] == '.'
-    && CreateTime[26] == 'Z'
-  }
 }
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy
index 763e4e9e5c..c277add48d 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminCreateKeys.dfy
@@ -4,8 +4,7 @@
 include "../src/Index.dfy"
 include "../../AwsCryptographyKeyStore/test/CleanupItems.dfy"
 include "../../AwsCryptographyKeyStore/test/Fixtures.dfy"
-include "../../AwsCryptographyKeyStore/test/TestGetKeys.dfy"
-include "../../AwsCryptographyKeyStore/Model/AwsCryptographyKeyStoreTypes.dfy"
+include "../../AwsCryptographyKeyStore/test/BranchKeyValidators.dfy"
 include "AdminFixtures.dfy"
 
 module {:options "/functionSyntax:4" } TestAdminCreateKeys {
@@ -22,9 +21,12 @@ module {:options "/functionSyntax:4" } TestAdminCreateKeys {
   import UUID
   import CleanupItems
   import AdminFixtures
-  import TestGetKeys
+  import BranchKeyValidators
 
-  method {:test} TestCreateBranchAndBeaconKeys()
+  // TODO-HV-2-FOLLOW : TestCreateSRKForHV1, TestCreateForHV1NoBKID, TestCreateForHV1BbkidNoECFails
+
+  const happyCaseId := "test-happy-case-admin-create-key-hv-1"
+  method {:test} TestCreateMRKForHV1()
   {
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
@@ -32,28 +34,145 @@ module {:options "/functionSyntax:4" } TestAdminCreateKeys {
     var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
     var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
     var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+    var keyStoreEast :- expect KeyStoreFromKMSConfig(
+      kmsConfig := KmsMrkConfigEast,
+      ddbClient? := Some(ddbClient)
+    );
+    var uuid :- expect UUID.GenerateUUID();
+    var bkid := happyCaseId + "-west-mrk-" + uuid;
+    // Create an HV-1 BK with a West MRK by calling us-west-2
+    var bk :- expect underTest.CreateKey(
+      Types.CreateKeyInput(
+        Identifier := Some(bkid),
+        EncryptionContext := Some(KmsMrkEC),
+        KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnWest),
+        // us-west-2 b/c CI region is us-west-2 and KMS Client is created
+        Strategy := Some(strategy)
+      ));
+    expect bk.Identifier == bkid;
+    // Validate an HV-1 BK with a West MRK by calling us-east-1
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStoreEast, storage,
+                                      encryptionContext := KmsMrkEC);
+    var keyStoreWest :- expect KeyStoreFromKMSConfig(
+      // Passing in KmsMrkConfigWest causes the KMS Client to be created for us-west-2
+      // bk's KMS-ARN is replicated to us-west-2
+      kmsConfig := KmsMrkConfigWest,
+      ddbClient? := Some(ddbClient)
+    );
+    // Validate an HV-1 BK with a West MRK by calling us-west-2
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStoreWest, storage,
+                                      encryptionContext := KmsMrkEC);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=bkid, ddbClient:=ddbClient);
+  }
 
-    var input := Types.CreateKeyInput(
-      Identifier := None,
-      EncryptionContext := None,
-      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
-      Strategy := Some(strategy)
+  // TODO-HV-2-M2?: Document that BKSA ONLY calls KMS/DDB for the region of the supplied clients.
+  // or Ensure BKSA has similar behavior as BKS in terms of MRK region.
+  const happyCaseIdHV2 := "test-happy-case-admin-create-key-hv-2"
+  method {:test} TestCreateMRKForHV2()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+    var keyStoreEast :- expect KeyStoreFromKMSConfig(
+      kmsConfig := KmsMrkConfigEast,
+      ddbClient? := Some(ddbClient)
     );
-    var identifier? :- expect underTest.CreateKey(input);
-    var identifier := identifier?.Identifier;
-    expect identifier?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v1,
-      "KeyStoreAdmin should create branch key with `hierarchy-version-1` when no `HierarchyVersion` provided";
-
-    // Get branch key items from storage
-    TestGetKeys.VerifyGetKeys(
-      identifier := identifier,
-      keyStore := keyStore,
-      storage := storage
+    var uuid :- expect UUID.GenerateUUID();
+    var bkid := happyCaseIdHV2 + "-west-mrk-" + uuid;
+    // Create an HV-2 BK with a West MRK by calling us-west-2
+    var bk :- expect underTest.CreateKey(
+      Types.CreateKeyInput(
+        Identifier := Some(bkid),
+        EncryptionContext := Some(RobbieEC),
+        KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnWest),
+        // us-west-2 b/c CI region is us-west-2 and KMS Client is created
+        Strategy := Some(strategy),
+        HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+      ));
+    expect bk.Identifier == bkid;
+    // Validate an HV-2 BK with a West MRK by calling us-east-1
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStoreEast, storage,
+                                      encryptionContext := RobbieEC,
+                                      hierarchyVersion := KeyStoreTypes.HierarchyVersion.v2);
+    var keyStoreWest :- expect KeyStoreFromKMSConfig(
+      // Passing in KmsMrkConfigWest causes the KMS Client to be created for us-west-2
+      // bk's KMS-ARN is replicated to us-west-2
+      kmsConfig := KmsMrkConfigWest,
+      ddbClient? := Some(ddbClient)
     );
+    // Validate an HV-2 BK with a West MRK by calling us-west-2
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStoreWest, storage,
+                                      encryptionContext := RobbieEC,
+                                      hierarchyVersion := KeyStoreTypes.HierarchyVersion.v2);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=bkid, ddbClient:=ddbClient);
+  }
 
-    // Since this process uses a read DDB table,
-    // the number of records will forever increase.
-    // To avoid this, remove the items.
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=identifier, ddbClient:=ddbClient);
+  method {:test} TestCreateSRKForHV2()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+    var uuid :- expect UUID.GenerateUUID();
+    var bkid := happyCaseIdHV2 + "-west-srk-" + uuid;
+    // Create an HV-2 BK with a West SRK by calling us-west-2
+    var bk :- expect underTest.CreateKey(
+      Types.CreateKeyInput(
+        Identifier := Some(bkid),
+        EncryptionContext := Some(RobbieEC),
+        KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+        // us-west-2 b/c CI region is us-west-2 and KMS Client is created
+        Strategy := Some(strategy),
+        HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+      ));
+    expect bk.Identifier == bkid;
+    // Validate an HV-2 BK with a West SRK by calling us-west-2
+    BranchKeyValidators.VerifyGetKeys(bkid, keyStore, storage,
+                                      encryptionContext := RobbieEC,
+                                      hierarchyVersion := KeyStoreTypes.HierarchyVersion.v2);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=bkid, ddbClient:=ddbClient);
+  }
+
+  method {:test} TestCreateSRKForHV2NoEC()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+    var bk :- expect underTest.CreateKey(
+      Types.CreateKeyInput(
+        KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+        Strategy := Some(strategy),
+        HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+      ));
+    BranchKeyValidators.VerifyGetKeys(bk.Identifier, keyStore, storage,
+                                      hierarchyVersion := KeyStoreTypes.HierarchyVersion.v2);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=bk.Identifier, ddbClient:=ddbClient);
+  }
+
+  method {:test} TestCreateSRKForHV2BkidWithNoECFails()
+  {
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
+    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+    var uuid :- expect UUID.GenerateUUID();
+    var bkid := happyCaseIdHV2 + "-west-srk-no-ec-fails-" + uuid;
+    var bk? := underTest.CreateKey(
+      Types.CreateKeyInput(
+        Identifier := Some(bkid),
+        KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+        Strategy := Some(strategy),
+        HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+      ));
+    expect bk?.Failure?, "Admin CreateKey, with an Identifier, but no EC, MUST Fail";
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=bkid, ddbClient:=ddbClient);
   }
 }
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
new file mode 100644
index 0000000000..458825900b
--- /dev/null
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
@@ -0,0 +1,123 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+include "../src/Index.dfy"
+include "../../AwsCryptographyKeyStore/test/CleanupItems.dfy"
+include "../../AwsCryptographyKeyStore/test/Fixtures.dfy"
+include "../../AwsCryptographyKeyStore/test/TestGetKeys.dfy"
+include "AdminFixtures.dfy"
+
+// TODO-HV-2-M2 : remove HV-1 restriction for Mutations
+module {:options "/functionSyntax:4" } TestAdminHV1Only {
+  import Types = AwsCryptographyKeyStoreAdminTypes
+  import KeyStoreAdmin
+  import KeyStore
+  import KeyStoreTypes = AwsCryptographyKeyStoreTypes
+  import ComAmazonawsKmsTypes
+  import KMS = Com.Amazonaws.Kms
+  import DDB = Com.Amazonaws.Dynamodb
+  import DefaultKeyStorageInterface
+  import opened Wrappers
+  import opened Fixtures
+  import UUID
+  import CleanupItems
+  import AdminFixtures
+  import TestGetKeys
+
+  // TODO-HV-2-M2 : Probably make this a happy test?
+  const testMutateForHV1WithAwsKmsSimpleFailsCaseId := "dafny-initialize-mutation-hv-1-simpleKms-rejection"
+  method {:test} TestMutateForHV1WithAwsKmsSimpleFails()
+  {
+    var uuid :- expect UUID.GenerateUUID();
+    var testId := testMutateForHV1WithAwsKmsSimpleFailsCaseId + "-" + uuid;
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var underTest :- expect AdminFixtures.DefaultAdmin();
+    var simpleStrategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
+    Fixtures.CreateHappyCaseId(id:=testId);
+
+    var mutationsRequest := Types.Mutations(
+      TerminalKmsArn := Some(Fixtures.postalHornKeyArn),
+      TerminalHierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v1)
+    );
+    var initInput := Types.InitializeMutationInput(
+      Identifier := testId,
+      Mutations := mutationsRequest,
+      Strategy := Some(simpleStrategy),
+      SystemKey := systemKey,
+      DoNotVersion := Some(true));
+    var initializeOutput := underTest.InitializeMutation(initInput);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
+    expect initializeOutput.Failure?, "Should have failed to InitializeMutation for HV-1 with Simple.";
+  }
+
+  // TODO-HV-2-M2 : Probably make this a happy test?
+  const testMutateInitEncountersHV2FailsCaseId := "dafny-initialize-mutation-encounters-hv-2-rejection"
+  const logPrefix := "\n" + testMutateInitEncountersHV2FailsCaseId + " :: "
+  method {:test} TestMutateInitEncountersHV2FailsCaseId()
+  {
+    var uuid :- expect UUID.GenerateUUID();
+    var testId := testMutateInitEncountersHV2FailsCaseId + "-" + uuid;
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var underTest :- expect AdminFixtures.DefaultAdmin();
+    var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
+    Fixtures.CreateHappyCaseId(id:=testId);
+    var _ :- expect AdminFixtures.AddAttributeWithoutLibrary(
+      id:=testId,
+      keyValue:=AdminFixtures.KeyValue(key:=Structure.HIERARCHY_VERSION, value:=Structure.HIERARCHY_VERSION_VALUE_2),
+      alsoViolateBeacon? := true, ddbClient? := Some(ddbClient),
+      kmsClient?:=Some(kmsClient), violateReservedAttribute:=true);
+
+    // print logPrefix + "re-wrote the item to be described as HV-2 even though it's not" + "\n";
+
+    var mutationsRequest := Types.Mutations(TerminalKmsArn := Some(Fixtures.postalHornKeyArn));
+    var initInput := Types.InitializeMutationInput(
+      Identifier := testId,
+      Mutations := mutationsRequest,
+      Strategy := Some(strategy),
+      SystemKey := systemKey,
+      DoNotVersion := Some(true));
+    var initializeOutput := underTest.InitializeMutation(initInput);
+
+    // print logPrefix + "initializeOutput :: ", initializeOutput, "\n";
+
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
+    expect initializeOutput.Failure?, "Should have failed InitializeMutation when HV-2 encountered by InitMutation.";
+  }
+
+  // TODO-HV-2-M3 : Probably make this a happy test?
+  const testVersionKeyEncountersHV2FailsCaseId := "dafny-version-key-encounters-hv-2-rejection"
+  const logPrefixVersion := "\n" + testVersionKeyEncountersHV2FailsCaseId + " :: "
+  method {:test} TestVersionKeyEncountersHV2Fails()
+  {
+    var uuid :- expect UUID.GenerateUUID();
+    var testId := testVersionKeyEncountersHV2FailsCaseId + "-" + uuid;
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var underTest :- expect AdminFixtures.DefaultAdmin();
+    var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
+    Fixtures.CreateHappyCaseId(id:=testId);
+    var _ :- expect AdminFixtures.AddAttributeWithoutLibrary(
+      id:=testId,
+      keyValue:=AdminFixtures.KeyValue(key:=Structure.HIERARCHY_VERSION, value:=Structure.HIERARCHY_VERSION_VALUE_2),
+      alsoViolateBeacon? := true, ddbClient? := Some(ddbClient),
+      kmsClient?:=Some(kmsClient), violateReservedAttribute:=true);
+
+    // print logPrefixVersion + "re-wrote the item to be described as HV-2 even though it's not" + "\n";
+
+    var versionInput := Types.VersionKeyInput(
+      Identifier := testId,
+      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
+      Strategy := Some(strategy));
+    var actualOutput := underTest.VersionKey(versionInput);
+
+    // print logPrefixVersion + "actualOutput :: ", actualOutput, "\n";
+
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
+    expect actualOutput.Failure?, "Should have failed VersionKey when HV-2 encountered by VersionKey.";
+  }
+}
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
deleted file mode 100644
index fd1c1044e4..0000000000
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV2.dfy
+++ /dev/null
@@ -1,326 +0,0 @@
-// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-include "../src/Index.dfy"
-include "../../AwsCryptographyKeyStore/test/CleanupItems.dfy"
-include "../../AwsCryptographyKeyStore/test/Fixtures.dfy"
-include "../../AwsCryptographyKeyStore/test/TestGetKeys.dfy"
-include "../../AwsCryptographyKeyStore/Model/AwsCryptographyKeyStoreTypes.dfy"
-include "AdminFixtures.dfy"
-
-// TODO-HV-2-M2 : remove HV-1 restriction for Mutations
-// TODO-HV-2-M3 : rename TestAdminHV2 -> TestAdminCreateKeyHV2
-//    and move HappyCases for version & mutations to it's own test modules
-module {:options "/functionSyntax:4" } TestAdminHV2 {
-  import Types = AwsCryptographyKeyStoreAdminTypes
-  import KeyStoreAdmin
-  import KeyStore
-  import KeyStoreTypes = AwsCryptographyKeyStoreTypes
-  import ComAmazonawsKmsTypes
-  import KMS = Com.Amazonaws.Kms
-  import DDB = Com.Amazonaws.Dynamodb
-  import DefaultKeyStorageInterface
-  import opened Wrappers
-  import opened Fixtures
-  import UUID
-  import CleanupItems
-  import AdminFixtures
-  import TestGetKeys
-
-  method {:test} TestCreateKeyHV2HappyCase()
-  {
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
-    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
-    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
-
-    var input := Types.CreateKeyInput(
-      Identifier := None,
-      EncryptionContext := None,
-      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
-      Strategy := Some(strategy),
-      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
-    );
-
-    var createKeyOutput? :- expect underTest.CreateKey(input);
-    expect createKeyOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
-
-    // Get branch key items from storage
-    TestGetKeys.VerifyGetKeys(
-      identifier := createKeyOutput?.Identifier,
-      keyStore := keyStore,
-      storage := storage
-    );
-
-    // Since this process uses a read DDB table,
-    // the number of records will forever increase.
-    // To avoid this, remove the items.
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=createKeyOutput?.Identifier, ddbClient:=ddbClient);
-  }
-
-  const happyCaseId := "test-create-key-hv2-happy-case"
-  method {:test} TestCreateKeyHV2WithOptions() {
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
-    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
-    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
-
-    // Create key with Custom EC & Branch Key Identifier
-    var uuid :- expect UUID.GenerateUUID();
-    var branchKeyId := happyCaseId + "-" + uuid;
-
-    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
-
-    var input := Types.CreateKeyInput(
-      Identifier := Some(branchKeyId),
-      EncryptionContext := Some(customEC),
-      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
-      Strategy := Some(strategy),
-      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
-    );
-
-    var createKeyOutput? :- expect underTest.CreateKey(input);
-    expect branchKeyId == createKeyOutput?.Identifier;
-    expect createKeyOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
-
-    // Get branch key items from storage
-    TestGetKeys.VerifyGetKeys(
-      identifier := branchKeyId,
-      encryptionContext := customEC,
-      keyStore := keyStore,
-      storage := storage
-    );
-
-    // Since this process uses a read DDB table,
-    // the number of records will forever increase.
-    // To avoid this, remove the items.
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=branchKeyId, ddbClient:=ddbClient);
-  }
-
-  // TODO-HV-2-M2?: Document that BKSA ONLY calls KMS/DDB for the region of the supplied clients.
-  // or Ensure BKSA has similar behavior as BKS in terms of MRK region.
-  /*
-  method {:test} TestCreateMRKForHV2()
-  {
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
-    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
-    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
-
-    // Create key with Custom EC & Branch Key Identifier
-    var uuid :- expect UUID.GenerateUUID();
-    var branchKeyIdWest := happyCaseId + "-" + WestBranchKey + "-" + uuid;
-    var branchKeyIdEast := happyCaseId + "-" + EastBranchKey + "-" + uuid;
-
-    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
-
-    var westOutput? :- expect underTest.CreateKey(Types.CreateKeyInput(
-                                                    Identifier := Some(branchKeyIdWest),
-                                                    EncryptionContext := Some(customEC),
-                                                    KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnWest),
-                                                    Strategy := Some(strategy),
-                                                    HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
-                                                  ));
-    expect westOutput?.Identifier == branchKeyIdWest;
-    expect westOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
-    print branchKeyIdWest + "\n";
-
-    var eastOutput? :- expect underTest.CreateKey(Types.CreateKeyInput(
-                                                    Identifier := Some(branchKeyIdEast),
-                                                    EncryptionContext := Some(customEC),
-                                                    KmsArn := Types.KmsSymmetricKeyArn.KmsMRKeyArn(MrkArnEast),
-                                                    Strategy := Some(strategy),
-                                                    HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
-                                                  ));
-    expect eastOutput?.Identifier == branchKeyIdEast;
-    expect eastOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
-    print branchKeyIdEast + "\n";
-  }
-  */
-
-  const testCreateKeyWithBKIDandNoECFailsId := "create-key-hv2-bkid-no-ec-fail"
-  method {:test} TestCreateKeyBKIDNoECFails()
-  {
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
-
-    var uuid :- expect UUID.GenerateUUID();
-    var branchKeyId := testCreateKeyWithBKIDandNoECFailsId + "-" + uuid;
-    var input := Types.CreateKeyInput(
-      Identifier := Some(branchKeyId),
-      EncryptionContext := None,
-      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
-      Strategy := Some(strategy),
-      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
-    );
-
-    var createKeyOutput := underTest.CreateKey(input);
-    expect createKeyOutput.Failure?, "Expected a failure when provided with Branch Key Identifier with no Custom Encryption Context";
-  }
-
-  method {:test} TestCreateKeyCustomECOnly()
-  {
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var storage :- expect Fixtures.DefaultStorage(ddbClient?:=Some(ddbClient));
-    var keyStore :- expect Fixtures.DefaultKeyStore(ddbClient?:=Some(ddbClient), kmsClient?:=Some(kmsClient));
-    var strategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
-
-    var customEC := map[UTF8.EncodeAscii("Koda") := UTF8.EncodeAscii("Is a dog.")];
-
-    var input := Types.CreateKeyInput(
-      Identifier := None,
-      EncryptionContext := Some(customEC),
-      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
-      Strategy := Some(strategy),
-      HierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
-    );
-
-    var createKeyOutput? :- expect underTest.CreateKey(input);
-    expect createKeyOutput?.HierarchyVersion == KeyStoreTypes.HierarchyVersion.v2;
-
-    // Get branch key items from storage
-    TestGetKeys.VerifyGetKeys(
-      identifier := createKeyOutput?.Identifier,
-      encryptionContext := customEC,
-      keyStore := keyStore,
-      storage := storage
-    );
-
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=createKeyOutput?.Identifier, ddbClient:=ddbClient);
-  }
-
-  // TODO-HV-2-M2 : Probably make this a happy test?
-  const testMutateForHV2FailsCaseId := "dafny-initialize-mutation-hv-2-rejection"
-  method {:test} TestMutateForHV2Fails()
-  {
-    var uuid :- expect UUID.GenerateUUID();
-    var testId := testMutateForHV2FailsCaseId + "-" + uuid;
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var underTest :- expect AdminFixtures.DefaultAdmin();
-    var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
-    Fixtures.CreateHappyCaseId(id:=testId);
-
-    var mutationsRequest := Types.Mutations(
-      TerminalKmsArn := Some(Fixtures.postalHornKeyArn),
-      TerminalHierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
-    );
-    var initInput := Types.InitializeMutationInput(
-      Identifier := testId,
-      Mutations := mutationsRequest,
-      Strategy := Some(strategy),
-      SystemKey := systemKey,
-      DoNotVersion := Some(true));
-    var initializeOutput := underTest.InitializeMutation(initInput);
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
-    expect initializeOutput.Failure?, "Should have failed to InitializeMutation HV-2.";
-  }
-
-  // TODO-HV-2-M2 : Probably make this a happy test?
-  const testMutateForHV1WithAwsKmsSimpleFailsCaseId := "dafny-initialize-mutation-hv-1-simpleKms-rejection"
-  method {:test} TestMutateForHV1WithAwsKmsSimpleFails()
-  {
-    var uuid :- expect UUID.GenerateUUID();
-    var testId := testMutateForHV1WithAwsKmsSimpleFailsCaseId + "-" + uuid;
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var underTest :- expect AdminFixtures.DefaultAdmin();
-    var simpleStrategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
-    Fixtures.CreateHappyCaseId(id:=testId);
-
-    var mutationsRequest := Types.Mutations(
-      TerminalKmsArn := Some(Fixtures.postalHornKeyArn),
-      TerminalHierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v1)
-    );
-    var initInput := Types.InitializeMutationInput(
-      Identifier := testId,
-      Mutations := mutationsRequest,
-      Strategy := Some(simpleStrategy),
-      SystemKey := systemKey,
-      DoNotVersion := Some(true));
-    var initializeOutput := underTest.InitializeMutation(initInput);
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
-    expect initializeOutput.Failure?, "Should have failed to InitializeMutation for HV-1 with Simple.";
-  }
-
-  // TODO-HV-2-M2 : Probably make this a happy test?
-  const testMutateInitEncountersHV2FailsCaseId := "dafny-initialize-mutation-encounters-hv-2-rejection"
-  const logPrefix := "\n" + testMutateInitEncountersHV2FailsCaseId + " :: "
-  method {:test} TestMutateInitEncountersHV2FailsCaseId()
-  {
-    var uuid :- expect UUID.GenerateUUID();
-    var testId := testMutateInitEncountersHV2FailsCaseId + "-" + uuid;
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var underTest :- expect AdminFixtures.DefaultAdmin();
-    var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
-    Fixtures.CreateHappyCaseId(id:=testId);
-    var _ :- expect AdminFixtures.AddAttributeWithoutLibrary(
-      id:=testId,
-      keyValue:=AdminFixtures.KeyValue(key:=Structure.HIERARCHY_VERSION, value:=Structure.HIERARCHY_VERSION_VALUE_2),
-      alsoViolateBeacon? := true, ddbClient? := Some(ddbClient),
-      kmsClient?:=Some(kmsClient), violateReservedAttribute:=true);
-
-    print logPrefix + "re-wrote the item to be described as HV-2 even though it's not" + "\n";
-
-    var mutationsRequest := Types.Mutations(TerminalKmsArn := Some(Fixtures.postalHornKeyArn));
-    var initInput := Types.InitializeMutationInput(
-      Identifier := testId,
-      Mutations := mutationsRequest,
-      Strategy := Some(strategy),
-      SystemKey := systemKey,
-      DoNotVersion := Some(true));
-    var initializeOutput := underTest.InitializeMutation(initInput);
-
-    print logPrefix + "initializeOutput :: ", initializeOutput, "\n";
-
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
-    expect initializeOutput.Failure?, "Should have failed InitializeMutation when HV-2 encountered by InitMutation.";
-  }
-
-  // TODO-HV-2-M2 : Probably make this a happy test?
-  const testVersionKeyEncountersHV2FailsCaseId := "dafny-version-key-encounters-hv-2-rejection"
-  const logPrefixVersion := "\n" + testVersionKeyEncountersHV2FailsCaseId + " :: "
-  method {:test} TestVersionKeyEncountersHV2Fails()
-  {
-    var uuid :- expect UUID.GenerateUUID();
-    var testId := testVersionKeyEncountersHV2FailsCaseId + "-" + uuid;
-    var ddbClient :- expect Fixtures.ProvideDDBClient();
-    var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var underTest :- expect AdminFixtures.DefaultAdmin();
-    var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
-    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
-    Fixtures.CreateHappyCaseId(id:=testId);
-    var _ :- expect AdminFixtures.AddAttributeWithoutLibrary(
-      id:=testId,
-      keyValue:=AdminFixtures.KeyValue(key:=Structure.HIERARCHY_VERSION, value:=Structure.HIERARCHY_VERSION_VALUE_2),
-      alsoViolateBeacon? := true, ddbClient? := Some(ddbClient),
-      kmsClient?:=Some(kmsClient), violateReservedAttribute:=true);
-
-    print logPrefixVersion + "re-wrote the item to be described as HV-2 even though it's not" + "\n";
-
-    var versionInput := Types.VersionKeyInput(
-      Identifier := testId,
-      KmsArn := Types.KmsSymmetricKeyArn.KmsKeyArn(keyArn),
-      Strategy := Some(strategy));
-    var actualOutput := underTest.VersionKey(versionInput);
-
-    print logPrefixVersion + "actualOutput :: ", actualOutput, "\n";
-
-    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
-    expect actualOutput.Failure?, "Should have failed VersionKey when HV-2 encountered by VersionKey.";
-  }
-}

From d1f7355100d2a79c0a7549439dcf7353c7ddb657 Mon Sep 17 00:00:00 2001
From: texastony <5892063+texastony@users.noreply.github.com>
Date: Tue, 8 Apr 2025 08:29:01 -0700
Subject: [PATCH 13/13] test(dafny): refactor BKSA tests

---
 .../test/BranchKeyValidators.dfy              |  2 +-
 .../test/TestAdminHV1Only.dfy                 | 33 +++++++++++++++++--
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy
index 6ad5163cd4..886006a6d9 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/test/BranchKeyValidators.dfy
@@ -184,8 +184,8 @@ module {:options "/functionSyntax:4" } BranchKeyValidators {
     expect |beaconKeyResult.beaconKeyMaterials.beaconKey.value| == 32, "Lenght of Beacon Key Materials MUST be 32 bytes.";
     if (beaconKeyResult.beaconKeyMaterials.encryptionContext != encryptionContext) {
       print "Beacon Key's encryption context is incorrect. Expected: ", encryptionContext, " but got: ", beaconKeyResult.beaconKeyMaterials.encryptionContext;
+      expect false, "Beacon Key's encryption context is incorrect.";
     }
-    expect beaconKeyResult.beaconKeyMaterials.encryptionContext == encryptionContext, "Beacon Key's encryption context is incorrect.";
     expect isValidBeaconResult?(beaconKeyResult, branchKeyId, encryptionContext);
     return beaconKeyResult.beaconKeyMaterials;
   }
diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
index 458825900b..715b74d6bb 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStoreAdmin/test/TestAdminHV1Only.dfy
@@ -24,6 +24,33 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
   import AdminFixtures
   import TestGetKeys
 
+  const testMutateForHV2FailsCaseId := "dafny-initialize-mutation-hv-2-rejection"
+  method {:test} TestMutateForHV2Fails()
+  {
+    var uuid :- expect UUID.GenerateUUID();
+    var testId := testMutateForHV2FailsCaseId + "-" + uuid;
+    var ddbClient :- expect Fixtures.ProvideDDBClient();
+    var kmsClient :- expect Fixtures.ProvideKMSClient();
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
+    var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
+    var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
+    Fixtures.CreateHappyCaseId(id:=testId);
+
+    var mutationsRequest := Types.Mutations(
+      TerminalKmsArn := Some(Fixtures.postalHornKeyArn),
+      TerminalHierarchyVersion := Some(KeyStoreTypes.HierarchyVersion.v2)
+    );
+    var initInput := Types.InitializeMutationInput(
+      Identifier := testId,
+      Mutations := mutationsRequest,
+      Strategy := Some(strategy),
+      SystemKey := systemKey,
+      DoNotVersion := Some(true));
+    var initializeOutput := underTest.InitializeMutation(initInput);
+    var _ := CleanupItems.DeleteBranchKey(Identifier:=testId, ddbClient:=ddbClient);
+    expect initializeOutput.Failure?, "Should have failed to InitializeMutation HV-2.";
+  }
+
   // TODO-HV-2-M2 : Probably make this a happy test?
   const testMutateForHV1WithAwsKmsSimpleFailsCaseId := "dafny-initialize-mutation-hv-1-simpleKms-rejection"
   method {:test} TestMutateForHV1WithAwsKmsSimpleFails()
@@ -32,7 +59,7 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     var testId := testMutateForHV1WithAwsKmsSimpleFailsCaseId + "-" + uuid;
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var underTest :- expect AdminFixtures.DefaultAdmin();
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
     var simpleStrategy :- expect AdminFixtures.SimpleKeyManagerStrategy(kmsClient?:=Some(kmsClient));
     var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
     Fixtures.CreateHappyCaseId(id:=testId);
@@ -61,7 +88,7 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     var testId := testMutateInitEncountersHV2FailsCaseId + "-" + uuid;
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var underTest :- expect AdminFixtures.DefaultAdmin();
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
     var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
     var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
     Fixtures.CreateHappyCaseId(id:=testId);
@@ -97,7 +124,7 @@ module {:options "/functionSyntax:4" } TestAdminHV1Only {
     var testId := testVersionKeyEncountersHV2FailsCaseId + "-" + uuid;
     var ddbClient :- expect Fixtures.ProvideDDBClient();
     var kmsClient :- expect Fixtures.ProvideKMSClient();
-    var underTest :- expect AdminFixtures.DefaultAdmin();
+    var underTest :- expect AdminFixtures.DefaultAdmin(ddbClient?:=Some(ddbClient));
     var strategy :- expect AdminFixtures.DefaultKeyManagerStrategy(kmsClient?:=Some(kmsClient));
     var systemKey := Types.SystemKey.trustStorage(trustStorage := Types.TrustStorage());
     Fixtures.CreateHappyCaseId(id:=testId);