Find archive URI for given account from ZK cache or ArchiveService

1 files changed, 24 insertions, 1 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 4bed7c5177f..7461d3aa47e 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
@@ -7,6 +7,7 @@ import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.tenant.ArchiveAccess;
 
 import java.net.URI;
+import java.time.Clock;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -19,9 +20,16 @@ import java.util.Set;
  */
 public class MockArchiveService implements ArchiveService {
 
+    private final Map<ZoneId, Set<TenantManagedArchiveBucket>> tenantArchiveBucketsByZone = new HashMap<>();
     public Set<VespaManagedArchiveBucket> archiveBuckets = new HashSet<>();
     public Map<TenantName, ArchiveAccess> authorizeAccessByTenantName = new HashMap<>();
 
+    private final Clock clock;
+
+    public MockArchiveService(Clock clock) {
+        this.clock = clock;
+    }
+
     @Override
     public VespaManagedArchiveBucket createArchiveBucketFor(ZoneId zoneId) {
         return new VespaManagedArchiveBucket("bucketName", "keyArn");
@@ -40,11 +48,26 @@ public class MockArchiveService implements ArchiveService {
 
     @Override
     public Optional<String> findEnclaveArchiveBucket(ZoneId zoneId, CloudAccount cloudAccount) {
-        return Optional.empty();
+        return tenantArchiveBucketsByZone.getOrDefault(zoneId, Set.of()).stream()
+                .filter(bucket -> bucket.cloudAccount().equals(cloudAccount))
+                .findFirst()
+                .map(TenantManagedArchiveBucket::bucketName);
     }
 
     @Override
     public URI bucketURI(ZoneId zoneId, String bucketName) {
         return URI.create(String.format("s3://%s/", bucketName));
     }
+
+
+    public void setEnclaveArchiveBucket(ZoneId zoneId, CloudAccount cloudAccount, String bucketName) {
+        removeEnclaveArchiveBucket(zoneId, cloudAccount);
+        tenantArchiveBucketsByZone.computeIfAbsent(zoneId, z -> new HashSet<>())
+                .add(new TenantManagedArchiveBucket(bucketName, cloudAccount, clock.instant()));
+    }
+
+    public void removeEnclaveArchiveBucket(ZoneId zoneId, CloudAccount cloudAccount) {
+        Optional.ofNullable(tenantArchiveBucketsByZone.get(zoneId))
+                .ifPresent(set -> set.removeIf(bucket -> bucket.cloudAccount().equals(cloudAccount)));
+    }
 }