1 files changed, 11 insertions, 3 deletions
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/archive/MockArchiveService.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/archive/MockArchiveService.java
index 5c979ddfc7b..796ca46ae4e 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/archive/MockArchiveService.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/archive/MockArchiveService.java
@@ -6,6 +6,8 @@ import com.yahoo.config.provision.zone.ZoneId;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
 
 /**
  * @author freva
@@ -13,7 +15,8 @@ import java.util.Map;
  */
 public class MockArchiveService implements ArchiveService {
 
-    public Map<ArchiveBucket, Map<TenantName, String>> authorizedIamRoles = new HashMap<>();
+    public Map<ArchiveBucket, Map<TenantName, String>> authorizedIamRolesForBucket = new HashMap<>();
+    public Map<String, Set<String>> authorizedIamRolesForKey = new TreeMap<>();
 
     @Override
     public ArchiveBucket createArchiveBucketFor(ZoneId zoneId, boolean sharded) {
@@ -21,7 +24,12 @@ public class MockArchiveService implements ArchiveService {
     }
 
     @Override
-    public void updateBucketAndKeyPolicy(ZoneId zoneId, ArchiveBucket bucket, Map<TenantName, String> authorizeIamRoleByTenantName) {
-        authorizedIamRoles.put(bucket, authorizeIamRoleByTenantName);
+    public void updateBucketPolicy(ZoneId zoneId, ArchiveBucket bucket, Map<TenantName, String> authorizeIamRoleByTenantName) {
+        authorizedIamRolesForBucket.put(bucket, authorizeIamRoleByTenantName);
+    }
+
+    @Override
+    public void updateKeyPolicy(ZoneId zoneId, String keyArn, Set<String> tenantAuthorizedIamRoles) {
+        authorizedIamRolesForKey.put(keyArn, tenantAuthorizedIamRoles);
     }
 }