fix tests

Author: RitvikKapila

DynamoDbEncryption/dafny/DynamoDbEncryption/test/BeaconTestFixtures.dfy

@@ -242,7 +242,9 @@ module BeaconTestFixtures {
       cache := MPT.Default(Default := MPT.DefaultCache(entryCapacity := 3))
     );
     var cache :- expect mpl.CreateCryptographicMaterialsCache(input);
-    return SI.KeySource(client, version.keyStore, SI.LiteralLoc(keys), cache, 0);
+    // Create a test partitionIdBytes
+    var partitionIdBytes : seq<uint8> := [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+    return SI.KeySource(client, version.keyStore, SI.LiteralLoc(keys), cache, 0, partitionIdBytes);
   }
 
   method GetMultiSource(keyName : string, version : BeaconVersion) returns (output : SI.KeySource)
@@ -257,7 +259,9 @@ module BeaconTestFixtures {
       cache := MPT.Default(Default := MPT.DefaultCache(entryCapacity := 3))
     );
     var cache :- expect mpl.CreateCryptographicMaterialsCache(input);
-    return SI.KeySource(client, version.keyStore, SI.MultiLoc(keyName, false), cache, 0);
+    // Create a test partitionIdBytes
+    var partitionIdBytes : seq<uint8> := [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+    return SI.KeySource(client, version.keyStore, SI.MultiLoc(keyName, false), cache, 0, partitionIdBytes);
   }
 
   const SimpleItem : DDB.AttributeMap := map[

TestVectors/dafny/DDBEncryption/src/JsonConfig.dfy

@@ -524,7 +524,10 @@ module {:options "-functionSyntax:4"} JsonConfig {
     var cache :- expect mpl.CreateCryptographicMaterialsCache(input);
 
     var client :- expect Primitives.AtomicPrimitives();
-    var src := SI.KeySource(client, store, SI.SingleLoc("foo"), cache, 100 as uint32);
+
+    // Create a test partitionIdBytes
+    var partitionIdBytes : seq<uint8> := [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+    var src := SI.KeySource(client, store, SI.SingleLoc("foo"), cache, 100 as uint32, partitionIdBytes);
 
     var bv :- expect SI.MakeBeaconVersion(1, src, map[], map[], map[]);
     return Success(bv);