1 files changed, 24 insertions, 10 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClientMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClientMock.java
index 3cee9e1ca2c..a6bcc9bf0ed 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClientMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClientMock.java
@@ -1,4 +1,4 @@
-// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.api.integration.billing;
 
 import com.yahoo.config.provision.TenantName;
@@ -26,9 +26,10 @@ public class BillingDatabaseClientMock implements BillingDatabaseClient {
     private final Map<Bill.Id, List<Bill.LineItem>> lineItems = new HashMap<>();
     private final Map<TenantName, List<Bill.LineItem>> uncommittedLineItems = new HashMap<>();
 
-    private final Map<Bill.Id, Bill.StatusHistory> statuses = new HashMap<>();
+    private final Map<Bill.Id, StatusHistory> statuses = new HashMap<>();
     private final Map<Bill.Id, ZonedDateTime> startTimes = new HashMap<>();
     private final Map<Bill.Id, ZonedDateTime> endTimes = new HashMap<>();
+    private final Map<Bill.Id, String> exportedInvoiceIds = new HashMap<>();
 
     private final ZonedDateTime startTime = LocalDate.of(2020, 4, 1).atStartOfDay(ZoneId.of("UTC"));
     private final ZonedDateTime endTime = LocalDate.of(2020, 5, 1).atStartOfDay(ZoneId.of("UTC"));
@@ -53,7 +54,7 @@ public class BillingDatabaseClientMock implements BillingDatabaseClient {
                 .findFirst();
     }
 
-    public String getStatus(Bill.Id invoiceId) {
+    public BillStatus getStatus(Bill.Id invoiceId) {
         return statuses.get(invoiceId).current();
     }
 
@@ -61,7 +62,7 @@ public class BillingDatabaseClientMock implements BillingDatabaseClient {
     public Bill.Id createBill(TenantName tenant, ZonedDateTime startTime, ZonedDateTime endTime, String agent) {
         var invoiceId = Bill.Id.generate();
         invoices.put(invoiceId, tenant);
-        statuses.computeIfAbsent(invoiceId, l -> Bill.StatusHistory.open(clock));
+        statuses.computeIfAbsent(invoiceId, l -> StatusHistory.open(clock));
         startTimes.put(invoiceId, startTime);
         endTimes.put(invoiceId, endTime);
         return invoiceId;
@@ -71,10 +72,11 @@ public class BillingDatabaseClientMock implements BillingDatabaseClient {
     public Optional<Bill> readBill(Bill.Id billId) {
         var invoice = Optional.ofNullable(invoices.get(billId));
         var lines = lineItems.getOrDefault(billId, List.of());
-        var status = statuses.getOrDefault(billId, Bill.StatusHistory.open(clock));
+        var status = statuses.getOrDefault(billId, StatusHistory.open(clock));
         var start = startTimes.getOrDefault(billId, startTime);
         var end = endTimes.getOrDefault(billId, endTime);
-        return invoice.map(tenant -> new Bill(billId, tenant, status, lines, start, end));
+        var exportedId = exportedInvoiceId(billId);
+        return invoice.map(tenant -> new Bill(billId, tenant, status, lines, start, end, exportedId));
     }
 
     @Override
@@ -88,8 +90,8 @@ public class BillingDatabaseClientMock implements BillingDatabaseClient {
     }
 
     @Override
-    public void setStatus(Bill.Id invoiceId, String agent, String status) {
-        statuses.computeIfAbsent(invoiceId, k -> Bill.StatusHistory.open(clock))
+    public void setStatus(Bill.Id invoiceId, String agent, BillStatus status) {
+        statuses.computeIfAbsent(invoiceId, k -> StatusHistory.open(clock))
                 .getHistory()
                 .put(ZonedDateTime.now(), status);
     }
@@ -157,7 +159,7 @@ public class BillingDatabaseClientMock implements BillingDatabaseClient {
                     var status = statuses.get(invoiceId);
                     var start = startTimes.get(invoiceId);
                     var end = endTimes.get(invoiceId);
-                    return new Bill(invoiceId, tenant, status, items, start, end);
+                    return new Bill(invoiceId, tenant, status, items, start, end, exportedInvoiceId(invoiceId));
                 })
                 .toList();
     }
@@ -171,11 +173,23 @@ public class BillingDatabaseClientMock implements BillingDatabaseClient {
                     var status = statuses.get(invoiceId);
                     var start = startTimes.get(invoiceId);
                     var end = endTimes.get(invoiceId);
-                    return new Bill(invoiceId, tenant, status, items, start, end);
+                    return new Bill(invoiceId, tenant, status, items, start, end, exportedInvoiceId(invoiceId));
                 })
                 .toList();
     }
 
     @Override
     public void maintain() {}
+
+    @Override
+    public void setExportedInvoiceId(Bill.Id billId, String invoiceId) {
+        exportedInvoiceIds.put(billId, invoiceId);
+    }
+
+    @Override
+    public void setExportedInvoiceItemId(String lineItemId, String invoiceItemId) { }
+
+    private String exportedInvoiceId(Bill.Id billId) {
+        return exportedInvoiceIds.getOrDefault(billId, null);
+    }
 }