killbill-aplcache

beatrix, usage: Enhance TestWithUsagePlugin and fix bugs

1/24/2019 9:15:42 PM

Details

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithUsagePlugin.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithUsagePlugin.java
index 687005d..f41359b 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithUsagePlugin.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithUsagePlugin.java
@@ -39,10 +39,8 @@ import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceItemType;
 import org.killbill.billing.osgi.api.OSGIServiceDescriptor;
 import org.killbill.billing.osgi.api.OSGIServiceRegistration;
-import org.killbill.billing.usage.api.SubscriptionUsageRecord;
-import org.killbill.billing.usage.api.UnitUsageRecord;
+import org.killbill.billing.usage.api.RolledUpUsage;
 import org.killbill.billing.usage.api.UsageApiException;
-import org.killbill.billing.usage.api.UsageRecord;
 import org.killbill.billing.usage.api.svcs.DefaultRawUsage;
 import org.killbill.billing.usage.plugin.api.RawUsageRecord;
 import org.killbill.billing.usage.plugin.api.UsagePluginApi;
@@ -54,6 +52,8 @@ import org.testng.annotations.Test;
 
 import com.google.common.collect.ImmutableSet;
 
+import static org.testng.Assert.assertEquals;
+
 public class TestWithUsagePlugin extends TestIntegrationBase {
 
     @Inject
@@ -118,6 +118,12 @@ public class TestWithUsagePlugin extends TestIntegrationBase {
         testUsagePluginApi.recordUsageData(aoSubscription.getId(), "tracking-1", "bullets", new LocalDate(2012, 4, 1), 99L, callContext);
         testUsagePluginApi.recordUsageData(aoSubscription.getId(), "tracking-2", "bullets", new LocalDate(2012, 4, 15), 100L, callContext);
 
+        // Wrong subscription - should be ignored...
+        testUsagePluginApi.recordUsageData(UUID.randomUUID(), "tracking-3", "bullets", new LocalDate(2012, 4, 5), 100L, callContext);
+
+        // Wrong unit - should be ignored...
+        testUsagePluginApi.recordUsageData(aoSubscription.getId(), "tracking-3", "bullets2", new LocalDate(2012, 4, 6), 200L, callContext);
+
         busHandler.pushExpectedEvents(NextEvent.PHASE, NextEvent.NULL_INVOICE, NextEvent.INVOICE, NextEvent.PAYMENT, NextEvent.INVOICE_PAYMENT);
         clock.addDays(30);
         assertListenerStatus();
@@ -127,11 +133,54 @@ public class TestWithUsagePlugin extends TestIntegrationBase {
                                                          new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.USAGE, new BigDecimal("5.90")));
         invoiceChecker.checkTrackingIds(curInvoice, ImmutableSet.of("tracking-1", "tracking-2"), internalCallContext);
 
+        final RolledUpUsage result1 = usageUserApi.getUsageForSubscription(aoSubscription.getId(), "bullets", new LocalDate(2012, 4, 1), new LocalDate(2012, 4, 15), callContext);
+        assertEquals(result1.getSubscriptionId(), aoSubscription.getId());
+        assertEquals(result1.getRolledUpUnits().size(), 1);
+        assertEquals(result1.getRolledUpUnits().get(0).getUnitType(), "bullets");
+        assertEquals(result1.getRolledUpUnits().get(0).getAmount().longValue(), 99L);
+
+        final List<LocalDate> transitionDates = new ArrayList<>();
+        transitionDates.add(new LocalDate(2012, 4, 1));
+        transitionDates.add(new LocalDate(2012, 4, 5));
+        transitionDates.add(new LocalDate(2012, 4, 6));
+        transitionDates.add(new LocalDate(2012, 4, 15));
+        transitionDates.add(new LocalDate(2012, 4, 17));
+
+        final List<RolledUpUsage> result2 = usageUserApi.getAllUsageForSubscription(aoSubscription.getId(), transitionDates, callContext);
+        assertEquals(result2.size(), 4);
+
+        assertEquals(result2.get(0).getSubscriptionId(), aoSubscription.getId());
+        assertEquals(result2.get(0).getStart(), new LocalDate(2012, 4, 1));
+        assertEquals(result2.get(0).getEnd(), new LocalDate(2012, 4, 5));
+        assertEquals(result2.get(0).getRolledUpUnits().size(), 1);
+        assertEquals(result2.get(0).getRolledUpUnits().get(0).getUnitType(), "bullets");
+        assertEquals(result2.get(0).getRolledUpUnits().get(0).getAmount().longValue(), 99L);
+
+        // Usage was for wrong subscriptionId
+        assertEquals(result2.get(1).getSubscriptionId(), aoSubscription.getId());
+        assertEquals(result2.get(1).getStart(), new LocalDate(2012, 4, 5));
+        assertEquals(result2.get(1).getEnd(), new LocalDate(2012, 4, 6));
+        assertEquals(result2.get(1).getRolledUpUnits().size(), 0);
+
+        // We see the wrong bullet2 unit
+        assertEquals(result2.get(2).getSubscriptionId(), aoSubscription.getId());
+        assertEquals(result2.get(2).getStart(), new LocalDate(2012, 4, 6));
+        assertEquals(result2.get(2).getEnd(), new LocalDate(2012, 4, 15));
+        assertEquals(result2.get(2).getRolledUpUnits().size(), 1);
+        assertEquals(result2.get(2).getRolledUpUnits().get(0).getUnitType(), "bullets2");
+        assertEquals(result2.get(2).getRolledUpUnits().get(0).getAmount().longValue(), 200L);
+
+        assertEquals(result2.get(3).getSubscriptionId(), aoSubscription.getId());
+        assertEquals(result2.get(3).getStart(), new LocalDate(2012, 4, 15));
+        assertEquals(result2.get(3).getEnd(), new LocalDate(2012, 4, 17));
+        assertEquals(result2.get(3).getRolledUpUnits().size(), 1);
+        assertEquals(result2.get(3).getRolledUpUnits().get(0).getUnitType(), "bullets");
+        assertEquals(result2.get(3).getRolledUpUnits().get(0).getAmount().longValue(), 100L);
+
     }
 
     public static class TestUsagePluginApi implements UsagePluginApi {
 
-
         private final SortedMap<LocalDate, List<RawUsageRecord>> usageData;
 
         public TestUsagePluginApi() {
@@ -155,7 +204,6 @@ public class TestWithUsagePlugin extends TestIntegrationBase {
 
         public void recordUsageData(final UUID subscriptionId, final String trackingId, final String unitType, final LocalDate startDate, final Long amount, final CallContext context) throws UsageApiException {
 
-
             List<RawUsageRecord> record = usageData.get(startDate);
             if (record == null) {
                 record = new LinkedList<>();
diff --git a/usage/src/main/java/org/killbill/billing/usage/api/BaseUserApi.java b/usage/src/main/java/org/killbill/billing/usage/api/BaseUserApi.java
index 7e4bca1..46ce9dd 100644
--- a/usage/src/main/java/org/killbill/billing/usage/api/BaseUserApi.java
+++ b/usage/src/main/java/org/killbill/billing/usage/api/BaseUserApi.java
@@ -25,15 +25,24 @@ import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.usage.plugin.api.RawUsageRecord;
 import org.killbill.billing.usage.plugin.api.UsagePluginApi;
 import org.killbill.billing.util.callcontext.TenantContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class BaseUserApi {
 
+    private static final Logger logger = LoggerFactory.getLogger(BaseUserApi.class);
+
     private final OSGIServiceRegistration<UsagePluginApi> pluginRegistry;
 
     public BaseUserApi(final OSGIServiceRegistration<UsagePluginApi> pluginRegistry) {
         this.pluginRegistry = pluginRegistry;
     }
 
+    //
+    // * If no plugin were registered or if plugins do not serve usage data for this tenant/account this returns null
+    //   and the usage module will look for data inside its own table.
+    // * If not, a possibly empty (or not) list should be returned (and the usage module will *not* look for data inside its own table)
+    //
     protected List<RawUsageRecord> getUsageFromPlugin(final LocalDate startDate, final LocalDate endDate, final TenantContext tenantContext) {
 
         final Set<String> allServices = pluginRegistry.getAllServices();
@@ -48,6 +57,13 @@ public class BaseUserApi {
             final List<RawUsageRecord> result = plugin.geUsageForAccount(startDate, endDate, tenantContext);
             // First plugin registered, returns result -- could be empty List if no usage was recorded.
             if (result != null) {
+
+                for (final RawUsageRecord cur : result) {
+                    if (cur.getDate().compareTo(startDate) < 0 || cur.getDate().compareTo(endDate) >=0) {
+                        logger.warn("Usage plugin returned usage data with date %s, not in the specified range [%s -> %s[",
+                                    cur.getDate(), startDate, endDate);
+                    }
+                }
                 return result;
             }
         }
diff --git a/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java b/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java
index 9f74a62..4128401 100644
--- a/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java
+++ b/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java
@@ -22,6 +22,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 import org.joda.time.LocalDate;
@@ -59,7 +60,7 @@ public class DefaultUsageUserApi extends BaseUserApi implements UsageUserApi {
                                final OSGIServiceRegistration<UsagePluginApi> pluginRegistry) {
         super(pluginRegistry);
         this.rolledUpUsageDao = rolledUpUsageDao;
-        this.internalCallContextFactory =internalCallContextFactory;
+        this.internalCallContextFactory = internalCallContextFactory;
     }
 
     @Override
@@ -85,11 +86,10 @@ public class DefaultUsageUserApi extends BaseUserApi implements UsageUserApi {
 
         final List<RawUsageRecord> rawUsage = getUsageFromPlugin(startDate, endDate, tenantContext);
         if (rawUsage != null) {
-            final List<RolledUpUnit> rolledUpAmount = getRolledUpUnitsForRawPluginUsage(rawUsage);
+            final List<RolledUpUnit> rolledUpAmount = getRolledUpUnitsForRawPluginUsage(subscriptionId, unitType, rawUsage);
             return new DefaultRolledUpUsage(subscriptionId, startDate, endDate, rolledUpAmount);
         }
 
-
         final List<RolledUpUsageModelDao> usageForSubscription = rolledUpUsageDao.getUsageForSubscription(subscriptionId, startDate, endDate, unitType, internalCallContextFactory.createInternalTenantContext(subscriptionId, ObjectType.SUBSCRIPTION, tenantContext));
         final List<RolledUpUnit> rolledUpAmount = getRolledUpUnits(usageForSubscription);
         return new DefaultRolledUpUsage(subscriptionId, startDate, endDate, rolledUpAmount);
@@ -106,7 +106,7 @@ public class DefaultUsageUserApi extends BaseUserApi implements UsageUserApi {
 
                 final List<RawUsageRecord> rawUsage = getUsageFromPlugin(prevDate, curDate, tenantContext);
                 if (rawUsage != null) {
-                    final List<RolledUpUnit> rolledUpAmount = getRolledUpUnitsForRawPluginUsage(rawUsage);
+                    final List<RolledUpUnit> rolledUpAmount = getRolledUpUnitsForRawPluginUsage(subscriptionId, null, rawUsage);
                     result.add(new DefaultRolledUpUsage(subscriptionId, prevDate, curDate, rolledUpAmount));
                 } else {
                     final List<RolledUpUsageModelDao> usageForSubscription = rolledUpUsageDao.getAllUsageForSubscription(subscriptionId, prevDate, curDate, internalCallContext);
@@ -119,23 +119,30 @@ public class DefaultUsageUserApi extends BaseUserApi implements UsageUserApi {
         return result;
     }
 
-
-    private List<RolledUpUnit> getRolledUpUnitsForRawPluginUsage(final List<RawUsageRecord> usageForSubscription) {
+    private List<RolledUpUnit> getRolledUpUnitsForRawPluginUsage(final UUID subscriptionId, @Nullable final String unitType, final List<RawUsageRecord> rawAccountUsage) {
         final Map<String, Long> tmp = new HashMap<String, Long>();
-        for (RawUsageRecord cur : usageForSubscription) {
+        for (RawUsageRecord cur : rawAccountUsage) {
+            // Filter out wrong subscriptionId
+            if (cur.getSubscriptionId().compareTo(subscriptionId) != 0) {
+                continue;
+            }
+
+            // Filter out wrong unitType if specified.
+            if (unitType != null && !unitType.equals(cur.getUnitType())) {
+                continue;
+            }
+
             Long currentAmount = tmp.get(cur.getUnitType());
             Long updatedAmount = (currentAmount != null) ? currentAmount + cur.getAmount() : cur.getAmount();
             tmp.put(cur.getUnitType(), updatedAmount);
         }
         final List<RolledUpUnit> result = new ArrayList<RolledUpUnit>(tmp.size());
-        for (final String unitType : tmp.keySet()) {
-            result.add(new DefaultRolledUpUnit(unitType, tmp.get(unitType)));
+        for (final String curType : tmp.keySet()) {
+            result.add(new DefaultRolledUpUnit(curType, tmp.get(curType)));
         }
         return result;
     }
 
-
-
     private List<RolledUpUnit> getRolledUpUnits(final List<RolledUpUsageModelDao> usageForSubscription) {
         final Map<String, Long> tmp = new HashMap<String, Long>();
         for (RolledUpUsageModelDao cur : usageForSubscription) {