killbill-aplcache

Add new policy to define how to compute usage price for consumable

3/17/2017 8:36:53 PM

Details

diff --git a/catalog/src/main/java/org/killbill/billing/catalog/CatalogSafetyInitializer.java b/catalog/src/main/java/org/killbill/billing/catalog/CatalogSafetyInitializer.java
index 60f56e3..c648b0e 100644
--- a/catalog/src/main/java/org/killbill/billing/catalog/CatalogSafetyInitializer.java
+++ b/catalog/src/main/java/org/killbill/billing/catalog/CatalogSafetyInitializer.java
@@ -28,6 +28,7 @@ import javax.xml.bind.annotation.XmlElementWrapper;
 
 import org.killbill.billing.catalog.api.BlockType;
 import org.killbill.billing.catalog.api.FixedType;
+import org.killbill.billing.catalog.api.TierBlockPolicy;
 
 public class CatalogSafetyInitializer {
 
@@ -58,6 +59,8 @@ public class CatalogSafetyInitializer {
                             initializeFieldWithValue(obj, f, FixedType.ONE_TIME);
                         } else if (BlockType.class.equals(f.getType())) {
                             initializeFieldWithValue(obj, f, BlockType.VANILLA);
+                        } else if (TierBlockPolicy.class.equals(f.getType())) {
+                            initializeFieldWithValue(obj, f, TierBlockPolicy.ALL_TIERS);
                         }
                     } else if (Integer.class.equals(f.getType())) {
                         initializeFieldWithValue(obj, f, DEFAULT_NON_REQUIRED_INTEGER_FIELD_VALUE);
@@ -97,6 +100,8 @@ public class CatalogSafetyInitializer {
                         result.add(f);
                     } else if (BlockType.class.equals(f.getType())) {
                         result.add(f);
+                    } else if (TierBlockPolicy.class.equals(f.getType())) {
+                        result.add(f);
                     }
                 } else if (Integer.class.equals(f.getType())) {
                     result.add(f);
diff --git a/catalog/src/main/java/org/killbill/billing/catalog/DefaultUsage.java b/catalog/src/main/java/org/killbill/billing/catalog/DefaultUsage.java
index 1619710..3786ce7 100644
--- a/catalog/src/main/java/org/killbill/billing/catalog/DefaultUsage.java
+++ b/catalog/src/main/java/org/killbill/billing/catalog/DefaultUsage.java
@@ -42,6 +42,8 @@ import org.killbill.billing.catalog.api.TieredBlockPriceOverride;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.catalog.api.UsagePriceOverride;
 import org.killbill.billing.catalog.api.UsageType;
+import org.killbill.billing.catalog.api.TierBlockPolicy;
+
 import org.killbill.xmlloader.ValidatingConfig;
 import org.killbill.xmlloader.ValidationError;
 import org.killbill.xmlloader.ValidationErrors;
@@ -63,6 +65,9 @@ public class DefaultUsage extends ValidatingConfig<StandaloneCatalog> implements
     @XmlAttribute(required = true)
     private UsageType usageType;
 
+    @XmlAttribute(required = false)
+    private TierBlockPolicy tierBlockPolicy;
+
     @XmlElement(required = true)
     private BillingPeriod billingPeriod;
 
@@ -98,6 +103,7 @@ public class DefaultUsage extends ValidatingConfig<StandaloneCatalog> implements
     public DefaultUsage(final Usage in, UsagePriceOverride override, Currency currency) {
               this.name = in.getName();
               this.usageType = in.getUsageType();
+              this.tierBlockPolicy = in.getTierBlockPolicy();
               this.billingPeriod = in.getBillingPeriod();
               this.billingMode = in.getBillingMode();
               this.limits = (DefaultLimit[]) in.getLimits();
@@ -154,6 +160,11 @@ public class DefaultUsage extends ValidatingConfig<StandaloneCatalog> implements
     }
 
     @Override
+    public TierBlockPolicy getTierBlockPolicy() {
+        return tierBlockPolicy;
+    }
+
+    @Override
     public BillingPeriod getBillingPeriod() {
         return billingPeriod;
     }
@@ -251,6 +262,11 @@ public class DefaultUsage extends ValidatingConfig<StandaloneCatalog> implements
         return this;
     }
 
+    public DefaultUsage setTierBlockPolicy(final TierBlockPolicy tierBlockPolicy) {
+        this.tierBlockPolicy = tierBlockPolicy;
+        return this;
+    }
+
     public DefaultUsage setPhase(final PlanPhase phase) {
         this.phase = phase;
         return this;
@@ -331,6 +347,9 @@ public class DefaultUsage extends ValidatingConfig<StandaloneCatalog> implements
         if (usageType != that.usageType) {
             return false;
         }
+        if (tierBlockPolicy != that.tierBlockPolicy) {
+            return false;
+        }
 
         return true;
     }
@@ -340,6 +359,7 @@ public class DefaultUsage extends ValidatingConfig<StandaloneCatalog> implements
         int result = name != null ? name.hashCode() : 0;
         result = 31 * result + (billingMode != null ? billingMode.hashCode() : 0);
         result = 31 * result + (usageType != null ? usageType.hashCode() : 0);
+        result = 31 * result + (tierBlockPolicy != null ? tierBlockPolicy.hashCode() : 0);
         result = 31 * result + (billingPeriod != null ? billingPeriod.hashCode() : 0);
         result = 31 * result + (limits != null ? Arrays.hashCode(limits) : 0);
         result = 31 * result + (blocks != null ? Arrays.hashCode(blocks) : 0);
diff --git a/catalog/src/test/java/org/killbill/billing/catalog/io/TestXMLReader.java b/catalog/src/test/java/org/killbill/billing/catalog/io/TestXMLReader.java
index ed4a141..2132d89 100644
--- a/catalog/src/test/java/org/killbill/billing/catalog/io/TestXMLReader.java
+++ b/catalog/src/test/java/org/killbill/billing/catalog/io/TestXMLReader.java
@@ -27,6 +27,7 @@ import org.killbill.billing.catalog.api.CatalogApiException;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.catalog.api.Plan;
 import org.killbill.billing.catalog.api.PlanPhase;
+import org.killbill.billing.catalog.api.TierBlockPolicy;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.catalog.api.UsageType;
 import org.killbill.xmlloader.XMLLoader;
@@ -71,6 +72,7 @@ public class TestXMLReader extends CatalogTestSuiteNoDB {
             assertEquals(usage.getBillingPeriod(), BillingPeriod.MONTHLY);
             assertEquals(usage.getUsageType(), UsageType.CAPACITY);
             assertEquals(usage.getBillingMode(), BillingMode.IN_ADVANCE);
+            assertEquals(usage.getTierBlockPolicy(), TierBlockPolicy.ALL_TIERS);
 
             assertEquals(usage.getBlocks().length, 0);
             assertEquals(usage.getTiers().length, 0);
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalUsageInArrear.java b/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalUsageInArrear.java
index 8d93aa4..8af590a 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalUsageInArrear.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalUsageInArrear.java
@@ -395,11 +395,24 @@ public class ContiguousIntervalUsageInArrear {
     BigDecimal computeToBeBilledConsumableInArrear(final RolledUpUnit roUnit) throws CatalogApiException {
 
         Preconditions.checkState(isBuilt.get());
+        final List<TieredBlock> tieredBlocks = getConsumableInArrearTieredBlocks(usage, roUnit.getUnitType());
 
-        BigDecimal result = BigDecimal.ZERO;
+        switch (usage.getTierBlockPolicy()) {
+            case ALL_TIERS:
+                return computeToBeBilledConsumableInArrearWith_ALL_TIERS(tieredBlocks, roUnit.getAmount());
+            case TOP_TIER:
+                return computeToBeBilledConsumableInArrearWith_TOP_TIER(tieredBlocks, roUnit.getAmount());
+            default:
+                throw new IllegalStateException("Unknwon TierBlockPolicy " + usage.getTierBlockPolicy());
 
-        final List<TieredBlock> tieredBlocks = getConsumableInArrearTieredBlocks(usage, roUnit.getUnitType());
-        int remainingUnits = roUnit.getAmount().intValue();
+        }
+    }
+
+
+    BigDecimal computeToBeBilledConsumableInArrearWith_ALL_TIERS(final List<TieredBlock> tieredBlocks, final Long units) throws CatalogApiException {
+
+        BigDecimal result = BigDecimal.ZERO;
+        int remainingUnits = units.intValue();
         for (final TieredBlock tieredBlock : tieredBlocks) {
 
             final int blockTierSize = tieredBlock.getSize().intValue();
@@ -417,6 +430,30 @@ public class ContiguousIntervalUsageInArrear {
         return result;
     }
 
+    BigDecimal computeToBeBilledConsumableInArrearWith_TOP_TIER(final List<TieredBlock> tieredBlocks, final Long units) throws CatalogApiException {
+
+        int remainingUnits = units.intValue();
+
+        // By default last last tierBlock
+        TieredBlock targetBlock = tieredBlocks.get(tieredBlocks.size() - 1);
+        // Loop through all tier block
+        for (final TieredBlock tieredBlock : tieredBlocks) {
+
+            final int blockTierSize = tieredBlock.getSize().intValue();
+            final int tmp = remainingUnits / blockTierSize + (remainingUnits % blockTierSize == 0 ? 0 : 1);
+            if (tmp > tieredBlock.getMax()) {
+                remainingUnits -= tieredBlock.getMax() * blockTierSize;
+            } else {
+                targetBlock = tieredBlock;
+                break;
+            }
+        }
+        final int lastBlockTierSize = targetBlock.getSize().intValue();
+        final int nbBlocks = units.intValue() / lastBlockTierSize + (units.intValue() % lastBlockTierSize == 0 ? 0 : 1);
+        return targetBlock.getPrice().getPrice(getCurrency()).multiply(new BigDecimal(nbBlocks));
+    }
+
+
     /**
      * @param filteredUsageForInterval the list of invoiceItem to consider
      * @return the price amount that was already billed for that period and usage section (across unitTypes)
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java
index ea6ce5e..5614cc0 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java
@@ -31,6 +31,7 @@ import org.killbill.billing.catalog.DefaultUsage;
 import org.killbill.billing.catalog.api.BillingPeriod;
 import org.killbill.billing.catalog.api.CatalogApiException;
 import org.killbill.billing.catalog.api.Currency;
+import org.killbill.billing.catalog.api.TierBlockPolicy;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.invoice.api.InvoiceItem;
 import org.killbill.billing.invoice.model.FixedPriceInvoiceItem;
@@ -84,7 +85,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTierWithBlocks(block);
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
 
         final LocalDate targetDate = startDate.plusDays(1);
         final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
@@ -116,21 +117,21 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
     }
 
     @Test(groups = "fast")
-    public void testComputeBilledUsage() throws CatalogApiException {
+    public void testComputeBilledUsageWith_ALL_TIERS() throws CatalogApiException {
 
         final DefaultTieredBlock block1 = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
         final DefaultTier tier1 = createDefaultTierWithBlocks(block1);
 
         final DefaultTieredBlock block2 = createDefaultTieredBlock("unit", 1000, 100, BigDecimal.ONE);
         final DefaultTier tier2 = createDefaultTierWithBlocks(block2);
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier1, tier2);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier1, tier2);
 
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
 
         final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
-                                                                                                                                                  BillingPeriod.MONTHLY,
-                                                                                                                                                  Collections.<Usage>emptyList())
+                                                                                                                                             BillingPeriod.MONTHLY,
+                                                                                                                                             Collections.<Usage>emptyList())
                                                                                                                      );
 
         final BigDecimal result = intervalConsumableInArrear.computeToBeBilledConsumableInArrear(new DefaultRolledUpUnit("unit", 5325L));
@@ -140,6 +141,51 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
     }
 
     @Test(groups = "fast")
+    public void testComputeBilledUsageWith_TOP_TIER() throws CatalogApiException {
+
+        final DefaultTieredBlock block1 = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
+        final DefaultTier tier1 = createDefaultTierWithBlocks(block1);
+
+        final DefaultTieredBlock block2 = createDefaultTieredBlock("unit", 1000, 100, BigDecimal.ONE);
+        final DefaultTier tier2 = createDefaultTierWithBlocks(block2);
+
+        final DefaultTieredBlock block3 = createDefaultTieredBlock("unit", 1000, 100, new BigDecimal("0.5"));
+        final DefaultTier tier3 = createDefaultTierWithBlocks(block3);
+
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.TOP_TIER, tier1, tier2, tier3);
+
+        final LocalDate targetDate = new LocalDate(2014, 03, 20);
+
+        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+                                                                                                                      createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
+                                                                                                                                             BillingPeriod.MONTHLY,
+                                                                                                                                             Collections.<Usage>emptyList())
+                                                                                                                     );
+        //
+        // In this model unit amount is first used to figure out which tier we are in, and then we price all unit at that 'target' tier
+        //
+        final BigDecimal inputTier1 = intervalConsumableInArrear.computeToBeBilledConsumableInArrear(new DefaultRolledUpUnit("unit", 1000L));
+        // 1000 units => (tier1) : 1000 / 100 + 1000 % 100 = 10
+        assertEquals(inputTier1, new BigDecimal("10"));
+
+        final BigDecimal inputTier2 = intervalConsumableInArrear.computeToBeBilledConsumableInArrear(new DefaultRolledUpUnit("unit", 101000L));
+        // 101000 units => (tier2) :  101000 / 1000 + 101000 % 1000 = 101 + 0 = 101
+        assertEquals(inputTier2, new BigDecimal("101"));
+
+        final BigDecimal inputTier3 = intervalConsumableInArrear.computeToBeBilledConsumableInArrear(new DefaultRolledUpUnit("unit", 101001L));
+        // 101001 units => (tier3) : 101001 / 1000 + 101001 % 1000 = 101 + 1 = 102 units => $51
+        assertEquals(inputTier3, new BigDecimal("51.0"));
+
+        // If we pass the maximum of the last tier, we price all units at the last tier
+        final BigDecimal inputLastTier = intervalConsumableInArrear.computeToBeBilledConsumableInArrear(new DefaultRolledUpUnit("unit", 300000L));
+        // 300000 units => (tier3) : 300000 / 1000 + 300000 % 1000 = 300 units => $150
+        assertEquals(inputLastTier, new BigDecimal("150.0"));
+
+    }
+
+
+
+    @Test(groups = "fast")
     public void testComputeMissingItems() throws CatalogApiException {
 
         final LocalDate startDate = new LocalDate(2014, 03, 20);
@@ -155,7 +201,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTierWithBlocks(block);
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
 
         final LocalDate targetDate = endDate;
 
@@ -216,7 +262,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
         final DefaultTier tier = createDefaultTierWithBlocks(tieredBlock1, tieredBlock2);
 
 
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
 
 
         final LocalDate t0 = new LocalDate(2015, 03, BCD);
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestRawUsageOptimizer.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestRawUsageOptimizer.java
index 38ab0f8..41201b9 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestRawUsageOptimizer.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestRawUsageOptimizer.java
@@ -29,6 +29,7 @@ import org.killbill.billing.catalog.DefaultTieredBlock;
 import org.killbill.billing.catalog.DefaultUsage;
 import org.killbill.billing.catalog.api.BillingPeriod;
 import org.killbill.billing.catalog.api.Currency;
+import org.killbill.billing.catalog.api.TierBlockPolicy;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.invoice.api.InvoiceItem;
 import org.killbill.billing.invoice.model.UsageInvoiceItem;
@@ -47,7 +48,7 @@ public class TestRawUsageOptimizer extends TestUsageInArrearBase {
         final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTierWithBlocks(block);
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
         knownUsage.put(usageName, usage);
 
         final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, firstEventStartDate.plusDays(1), invoiceItems, knownUsage, internalCallContext);
@@ -66,7 +67,7 @@ public class TestRawUsageOptimizer extends TestUsageInArrearBase {
         final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTierWithBlocks(block);
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
         knownUsage.put(usageName, usage);
 
         final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, targetDate, invoiceItems, knownUsage, internalCallContext);
@@ -89,7 +90,7 @@ public class TestRawUsageOptimizer extends TestUsageInArrearBase {
         final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTierWithBlocks(block);
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
         knownUsage.put(usageName, usage);
 
         final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, targetDate, invoiceItems, knownUsage, internalCallContext);
@@ -111,12 +112,12 @@ public class TestRawUsageOptimizer extends TestUsageInArrearBase {
         final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTierWithBlocks(block);
-        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, tier);
+        final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
         knownUsage.put(usageName, usage);
 
         final DefaultTieredBlock block2 = createDefaultTieredBlock("unit2", 10, 10000, BigDecimal.TEN);
         final DefaultTier tier2 = createDefaultTierWithBlocks(block2);
-        final DefaultUsage usage2 = createConsumableInArrearUsage("usageName2", BillingPeriod.ANNUAL, tier2);
+        final DefaultUsage usage2 = createConsumableInArrearUsage("usageName2", BillingPeriod.ANNUAL, TierBlockPolicy.ALL_TIERS, tier2);
         knownUsage.put("usageName2", usage2);
 
         final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, targetDate, invoiceItems, knownUsage, internalCallContext);
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java
index 3e6b03a..eb02634 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java
@@ -26,6 +26,7 @@ import org.joda.time.LocalDate;
 import org.killbill.billing.catalog.DefaultTier;
 import org.killbill.billing.catalog.DefaultTieredBlock;
 import org.killbill.billing.catalog.api.BillingPeriod;
+import org.killbill.billing.catalog.api.TierBlockPolicy;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.junction.BillingEvent;
 import org.killbill.billing.usage.RawUsage;
@@ -53,12 +54,12 @@ public class TestSubscriptionConsumableInArrear extends TestUsageInArrearBase {
         final String usageName1 = "erw";
         final DefaultTieredBlock block1 = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
         final DefaultTier tier1 = createDefaultTierWithBlocks(block1);
-        final Usage usage1 = createConsumableInArrearUsage(usageName1, BillingPeriod.MONTHLY, tier1);
+        final Usage usage1 = createConsumableInArrearUsage(usageName1, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier1);
 
         final String usageName2 = "hghg";
         final DefaultTieredBlock block2 = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
         final DefaultTier tier2 = createDefaultTierWithBlocks(block2);
-        final Usage usage2 = createConsumableInArrearUsage(usageName2, BillingPeriod.MONTHLY, tier2);
+        final Usage usage2 = createConsumableInArrearUsage(usageName2, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier2);
 
         final DateTime dt1 = new DateTime(2013, 3, 23, 4, 34, 59, DateTimeZone.UTC);
         final BillingEvent evt1 = createMockBillingEvent(dt1, BillingPeriod.MONTHLY, ImmutableList.<Usage>builder().add(usage1).add(usage2).build());
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java
index d766b3a..e8ba727 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java
@@ -37,6 +37,7 @@ import org.killbill.billing.catalog.api.BillingPeriod;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.catalog.api.Plan;
 import org.killbill.billing.catalog.api.PlanPhase;
+import org.killbill.billing.catalog.api.TierBlockPolicy;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.catalog.api.UsageType;
 import org.killbill.billing.invoice.InvoiceTestSuiteNoDB;
@@ -81,11 +82,12 @@ public abstract class TestUsageInArrearBase extends InvoiceTestSuiteNoDB {
         return intervalConsumableInArrear;
     }
 
-    protected DefaultUsage createConsumableInArrearUsage(final String usageName, final BillingPeriod billingPeriod, final DefaultTier... tiers) {
+    protected DefaultUsage createConsumableInArrearUsage(final String usageName, final BillingPeriod billingPeriod, final TierBlockPolicy tierBlockPolicy, final DefaultTier... tiers) {
         final DefaultUsage usage = new DefaultUsage();
         usage.setName(usageName);
         usage.setBillingMode(BillingMode.IN_ARREAR);
         usage.setUsageType(UsageType.CONSUMABLE);
+        usage.setTierBlockPolicy(tierBlockPolicy);
         usage.setBillingPeriod(billingPeriod);
         usage.setTiers(tiers);
         return usage;