killbill-memoizeit

invoice: Refactor usage code to extract CONSUMABLE and CAPACITY

2/10/2018 3:09:02 AM

Details

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalCapacityUsageInArrear.java b/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalCapacityUsageInArrear.java
new file mode 100644
index 0000000..02c116b
--- /dev/null
+++ b/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalCapacityUsageInArrear.java
@@ -0,0 +1,138 @@
+/*
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.killbill.billing.invoice.usage;
+
+import java.math.BigDecimal;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.UUID;
+
+import org.joda.time.LocalDate;
+import org.killbill.billing.callcontext.InternalTenantContext;
+import org.killbill.billing.catalog.api.CatalogApiException;
+import org.killbill.billing.catalog.api.Limit;
+import org.killbill.billing.catalog.api.Tier;
+import org.killbill.billing.catalog.api.Usage;
+import org.killbill.billing.invoice.api.InvoiceItem;
+import org.killbill.billing.invoice.model.UsageInvoiceItem;
+import org.killbill.billing.invoice.usage.details.UsageCapacityInArrearDetail;
+import org.killbill.billing.invoice.usage.details.UsageConsumableInArrearTierUnitDetail;
+import org.killbill.billing.invoice.usage.details.UsageInArrearDetail;
+import org.killbill.billing.invoice.usage.details.UsageInArrearTierUnitDetail;
+import org.killbill.billing.usage.RawUsage;
+import org.killbill.billing.usage.api.RolledUpUnit;
+import org.killbill.billing.util.config.definition.InvoiceConfig.UsageDetailMode;
+import org.killbill.billing.util.jackson.ObjectMapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Joiner;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Lists;
+
+import static org.killbill.billing.invoice.usage.UsageUtils.getCapacityInArrearTier;
+
+public class ContiguousIntervalCapacityUsageInArrear extends ContiguousIntervalUsageInArrear {
+
+    private static final Logger log = LoggerFactory.getLogger(ContiguousIntervalCapacityUsageInArrear.class);
+
+    private static final ObjectMapper objectMapper = new ObjectMapper();
+
+    public ContiguousIntervalCapacityUsageInArrear(final Usage usage,
+                                                   final UUID accountId,
+                                                   final UUID invoiceId,
+                                                   final List<RawUsage> rawSubscriptionUsage,
+                                                   final LocalDate targetDate,
+                                                   final LocalDate rawUsageStartDate,
+                                                   final UsageDetailMode usageDetailMode,
+                                                   final InternalTenantContext internalTenantContext) {
+        super(usage, accountId, invoiceId, rawSubscriptionUsage, targetDate, rawUsageStartDate, usageDetailMode, internalTenantContext);
+    }
+
+
+
+    @Override
+    protected void populateResults(final LocalDate startDate, final LocalDate endDate, final Iterable<InvoiceItem> billedItems, final BigDecimal billedUsage, final BigDecimal toBeBilledUsage, final UsageInArrearDetail toBeBilledUsageDetails, final boolean areAllBilledItemsWithDetails, final List<InvoiceItem> result) {
+        // Compute final amount by subtracting  amount that was already billed.
+        if (!billedItems.iterator().hasNext() || billedUsage.compareTo(toBeBilledUsage) < 0) {
+            final BigDecimal amountToBill = toBeBilledUsage.subtract(billedUsage);
+
+            if (amountToBill.compareTo(BigDecimal.ZERO) > 0) {
+                    final String itemDetails = areAllBilledItemsWithDetails ? toBeBilledUsageDetails.toJson(objectMapper) : null;
+                    final InvoiceItem item = new UsageInvoiceItem(invoiceId, accountId, getBundleId(), getSubscriptionId(), getPlanName(),
+                                                                  getPhaseName(), usage.getName(), startDate, endDate, amountToBill, null, getCurrency(), null, itemDetails);
+                    result.add(item);
+            }
+        }
+
+    }
+
+
+    @Override
+    protected UsageInArrearDetail getToBeBilledUsageDetails(final List<RolledUpUnit> rolledUpUnits, final Iterable<InvoiceItem> billedItems, final boolean areAllBilledItemsWithDetails) throws CatalogApiException {
+        return computeToBeBilledCapacityInArrear(rolledUpUnits);
+    }
+
+    private Limit getTierLimit(final Tier tier, final String unitType) {
+        for (final Limit cur : tier.getLimits()) {
+            if (cur.getUnit().getName().equals(unitType)) {
+                return cur;
+            }
+        }
+        Preconditions.checkState(false, "Could not find unit type " + unitType + " in usage tier ");
+        return null;
+    }
+
+    @VisibleForTesting
+    UsageCapacityInArrearDetail computeToBeBilledCapacityInArrear(final List<RolledUpUnit> roUnits) throws CatalogApiException {
+        Preconditions.checkState(isBuilt.get());
+
+        final List<Tier> tiers = getCapacityInArrearTier(usage);
+
+        final Set<String> perUnitTypeDetailTierLevel = new HashSet<String>();
+        int tierNum = 0;
+        final List<UsageInArrearTierUnitDetail> toBeBilledDetails = Lists.newLinkedList();
+        for (final Tier cur : tiers) {
+            tierNum++;
+            boolean complies = true;
+            for (final RolledUpUnit ro : roUnits) {
+                final Limit tierLimit = getTierLimit(cur, ro.getUnitType());
+                // We ignore the min and only look at the max Limit as the tiers should be contiguous.
+                // Specifying a -1 value for last max tier will make the validation works
+                if (tierLimit.getMax() != (double) -1 && ro.getAmount().doubleValue() > tierLimit.getMax()) {
+                    complies = false;
+                } else {
+                    if (!perUnitTypeDetailTierLevel.contains(ro.getUnitType())) {
+                        toBeBilledDetails.add(new UsageConsumableInArrearTierUnitDetail(tierNum, ro.getUnitType(), cur.getRecurringPrice().getPrice(getCurrency()), 1, ro.getAmount().intValue(), BigDecimal.ZERO));
+                        perUnitTypeDetailTierLevel.add(ro.getUnitType());
+                    }
+                }
+            }
+            if (complies) {
+                return new UsageCapacityInArrearDetail(toBeBilledDetails, cur.getRecurringPrice().getPrice(getCurrency()));
+            }
+        }
+        // Probably invalid catalog config
+        final Joiner joiner = Joiner.on(", ");
+        joiner.join(roUnits);
+        Preconditions.checkState(false, "Could not find tier for usage " + usage.getName() + "matching with data = " + joiner.join(roUnits));
+        return null;
+    }
+}
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalConsumableUsageInArrear.java b/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalConsumableUsageInArrear.java
new file mode 100644
index 0000000..b540981
--- /dev/null
+++ b/invoice/src/main/java/org/killbill/billing/invoice/usage/ContiguousIntervalConsumableUsageInArrear.java
@@ -0,0 +1,245 @@
+/*
+ * Copyright 2014-2016 Groupon, Inc
+ * Copyright 2014-2016 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.killbill.billing.invoice.usage;
+
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import org.joda.time.LocalDate;
+import org.killbill.billing.callcontext.InternalTenantContext;
+import org.killbill.billing.catalog.api.CatalogApiException;
+import org.killbill.billing.catalog.api.TieredBlock;
+import org.killbill.billing.catalog.api.Usage;
+import org.killbill.billing.invoice.api.InvoiceItem;
+import org.killbill.billing.invoice.model.UsageInvoiceItem;
+import org.killbill.billing.invoice.usage.details.UsageConsumableInArrearDetail;
+import org.killbill.billing.invoice.usage.details.UsageConsumableInArrearTierUnitDetail;
+import org.killbill.billing.invoice.usage.details.UsageInArrearDetail;
+import org.killbill.billing.usage.RawUsage;
+import org.killbill.billing.usage.api.RolledUpUnit;
+import org.killbill.billing.util.config.definition.InvoiceConfig.UsageDetailMode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+
+import static org.killbill.billing.invoice.usage.UsageUtils.getConsumableInArrearTieredBlocks;
+
+public class ContiguousIntervalConsumableUsageInArrear extends ContiguousIntervalUsageInArrear {
+
+    private static final Logger log = LoggerFactory.getLogger(ContiguousIntervalConsumableUsageInArrear.class);
+
+    public ContiguousIntervalConsumableUsageInArrear(final Usage usage,
+                                                     final UUID accountId,
+                                                     final UUID invoiceId,
+                                                     final List<RawUsage> rawSubscriptionUsage,
+                                                     final LocalDate targetDate,
+                                                     final LocalDate rawUsageStartDate,
+                                                     final UsageDetailMode usageDetailMode,
+                                                     final InternalTenantContext internalTenantContext) {
+        super(usage, accountId, invoiceId, rawSubscriptionUsage, targetDate, rawUsageStartDate, usageDetailMode, internalTenantContext);
+    }
+
+    @Override
+    protected void populateResults(final LocalDate startDate, final LocalDate endDate, final Iterable<InvoiceItem> billedItems, final BigDecimal billedUsage, final BigDecimal toBeBilledUsage, final UsageInArrearDetail toBeBilledUsageDetails, final boolean areAllBilledItemsWithDetails, final List<InvoiceItem> result) {
+        // Compute final amount by subtracting  amount that was already billed.
+        if (!billedItems.iterator().hasNext() || billedUsage.compareTo(toBeBilledUsage) < 0) {
+            final BigDecimal amountToBill = areAllBilledItemsWithDetails ? toBeBilledUsage : toBeBilledUsage.subtract(billedUsage);
+
+            if (amountToBill.compareTo(BigDecimal.ZERO) > 0) {
+                if (UsageDetailMode.DETAIL == usageDetailMode) {
+                    for (UsageConsumableInArrearTierUnitDetail toBeBilledUsageDetail : ((UsageConsumableInArrearDetail) toBeBilledUsageDetails).getTierDetails()) {
+                        final InvoiceItem item = new UsageInvoiceItem(invoiceId, accountId, getBundleId(), getSubscriptionId(), getPlanName(),
+                                                                      getPhaseName(), usage.getName(), startDate, endDate, toBeBilledUsageDetail.getAmount(), toBeBilledUsageDetail.getTierPrice(), getCurrency(), toBeBilledUsageDetail.getQuantity(), null);
+                        result.add(item);
+                    }
+                } else {
+                    final String itemDetails = areAllBilledItemsWithDetails ? toBeBilledUsageDetails.toJson(objectMapper) : null;
+                    final InvoiceItem item = new UsageInvoiceItem(invoiceId, accountId, getBundleId(), getSubscriptionId(), getPlanName(),
+                                                                  getPhaseName(), usage.getName(), startDate, endDate, amountToBill, null, getCurrency(), null, itemDetails);
+                    result.add(item);
+                }
+            }
+        }
+
+    }
+
+    protected UsageInArrearDetail getToBeBilledUsageDetails(final List<RolledUpUnit> rolledUpUnits, final Iterable<InvoiceItem> billedItems, final boolean areAllBilledItemsWithDetails) throws CatalogApiException {
+
+        final Map<String, List<UsageConsumableInArrearTierUnitDetail>> previousUnitsUsage;
+        if (areAllBilledItemsWithDetails) {
+            previousUnitsUsage = new HashMap<String, List<UsageConsumableInArrearTierUnitDetail>>();
+            for (RolledUpUnit cur : rolledUpUnits) {
+                final List<UsageConsumableInArrearTierUnitDetail> usageInArrearDetailForUnitType = getUsageConsumableInArrearDetailForUnitType(billedItems, cur.getUnitType());
+                previousUnitsUsage.put(cur.getUnitType(), usageInArrearDetailForUnitType);
+            }
+        } else {
+            previousUnitsUsage = ImmutableMap.of();
+        }
+
+        final List<UsageConsumableInArrearTierUnitDetail> usageConsumableInArrearTierUnitDetails = new ArrayList<UsageConsumableInArrearTierUnitDetail>();
+        for (final RolledUpUnit cur : rolledUpUnits) {
+            if (!unitTypes.contains(cur.getUnitType())) {
+                log.warn("ContiguousIntervalConsumableInArrear is skipping unitType " + cur.getUnitType());
+                continue;
+            }
+            final List<UsageConsumableInArrearTierUnitDetail> previousUsage = previousUnitsUsage.containsKey(cur.getUnitType()) ? previousUnitsUsage.get(cur.getUnitType()) : ImmutableList.<UsageConsumableInArrearTierUnitDetail>of();
+
+            usageConsumableInArrearTierUnitDetails.addAll(computeToBeBilledConsumableInArrear(cur, previousUsage, areAllBilledItemsWithDetails));
+        }
+        final UsageInArrearDetail toBeBilledUsageDetails = new UsageConsumableInArrearDetail(usageConsumableInArrearTierUnitDetails);
+        return toBeBilledUsageDetails;
+    }
+
+    private List<UsageConsumableInArrearTierUnitDetail> getUsageConsumableInArrearDetailForUnitType(final Iterable<InvoiceItem> billedItems, final String unitType) {
+
+        final List<UsageConsumableInArrearTierUnitDetail> result = new ArrayList<UsageConsumableInArrearTierUnitDetail>();
+        for (final InvoiceItem bi : billedItems) {
+            final List<UsageConsumableInArrearTierUnitDetail> billedUsageItemDetails = fromJson(bi.getItemDetails());
+            for (final UsageConsumableInArrearTierUnitDetail curDetail : billedUsageItemDetails) {
+                if (curDetail.getTierUnit().equals(unitType)) {
+                    result.add(curDetail);
+                }
+            }
+        }
+        return result;
+    }
+
+    @VisibleForTesting
+    List<UsageConsumableInArrearTierUnitDetail> computeToBeBilledConsumableInArrear(final RolledUpUnit roUnit, final List<UsageConsumableInArrearTierUnitDetail> previousUsage, final boolean areAllBilledItemsWithDetails) throws CatalogApiException {
+
+        Preconditions.checkState(isBuilt.get());
+        final List<TieredBlock> tieredBlocks = getConsumableInArrearTieredBlocks(usage, roUnit.getUnitType());
+
+        switch (usage.getTierBlockPolicy()) {
+            case ALL_TIERS:
+                return computeToBeBilledConsumableInArrearWith_ALL_TIERS(tieredBlocks, previousUsage, roUnit.getAmount());
+            case TOP_TIER:
+                return Arrays.asList(computeToBeBilledConsumableInArrearWith_TOP_TIER(tieredBlocks, previousUsage, roUnit.getAmount(), areAllBilledItemsWithDetails));
+            default:
+                throw new IllegalStateException("Unknown TierBlockPolicy " + usage.getTierBlockPolicy());
+        }
+    }
+
+    List<UsageConsumableInArrearTierUnitDetail> computeToBeBilledConsumableInArrearWith_ALL_TIERS(final List<TieredBlock> tieredBlocks, final List<UsageConsumableInArrearTierUnitDetail> previousUsage, final Long units) throws CatalogApiException {
+
+        List<UsageConsumableInArrearTierUnitDetail> toBeBilledDetails = Lists.newLinkedList();
+        int remainingUnits = units.intValue();
+        int tierNum = 0;
+
+        final int lastPreviousUsageTier = previousUsage.size(); // we count tier from 1, 2, ...
+        final boolean hasPreviousUsage = lastPreviousUsageTier > 0;
+
+        for (final TieredBlock tieredBlock : tieredBlocks) {
+
+            tierNum++;
+            final int blockTierSize = tieredBlock.getSize().intValue();
+            final int tmp = remainingUnits / blockTierSize + (remainingUnits % blockTierSize == 0 ? 0 : 1);
+            int nbUsedTierBlocks;
+            if (tmp > tieredBlock.getMax()) {
+                nbUsedTierBlocks = tieredBlock.getMax().intValue();
+                remainingUnits -= tieredBlock.getMax() * blockTierSize;
+            } else {
+                nbUsedTierBlocks = tmp;
+                remainingUnits = 0;
+            }
+
+            if (nbUsedTierBlocks > 0) {
+                if (hasPreviousUsage) {
+                    final Integer previousUsageQuantity = tierNum <= lastPreviousUsageTier ? previousUsage.get(tierNum - 1).getQuantity() : 0;
+                    if (tierNum < lastPreviousUsageTier) {
+                        Preconditions.checkState(nbUsedTierBlocks == previousUsageQuantity, String.format("Expected usage for tier='%d', unit='%s' to be full, instead found units='[%d/%d]'",
+                                                                                                          tierNum, tieredBlock.getUnit().getName(), nbUsedTierBlocks, previousUsageQuantity));
+                    } else {
+                        Preconditions.checkState(nbUsedTierBlocks - previousUsageQuantity >= 0, String.format("Expected usage for tier='%d', unit='%s' to contain at least as mush as current usage, instead found units='[%d/%d]",
+                                                                                                              tierNum, tieredBlock.getUnit().getName(), nbUsedTierBlocks, previousUsageQuantity));
+                    }
+                    nbUsedTierBlocks = nbUsedTierBlocks - previousUsageQuantity;
+                }
+                if (nbUsedTierBlocks > 0) {
+                    toBeBilledDetails.add(new UsageConsumableInArrearTierUnitDetail(tierNum, tieredBlock.getUnit().getName(), tieredBlock.getPrice().getPrice(getCurrency()), blockTierSize, nbUsedTierBlocks));
+                }
+            }
+        }
+        return toBeBilledDetails;
+    }
+
+    UsageConsumableInArrearTierUnitDetail computeToBeBilledConsumableInArrearWith_TOP_TIER(final List<TieredBlock> tieredBlocks, final List<UsageConsumableInArrearTierUnitDetail> previousUsage, final Long units, final boolean areAllBilledItemsWithDetails) throws CatalogApiException {
+
+        int remainingUnits = units.intValue();
+
+        // By default last last tierBlock
+        TieredBlock targetBlock = tieredBlocks.get(tieredBlocks.size() - 1);
+        int targetTierNum = tieredBlocks.size();
+        int tierNum = 0;
+        // Loop through all tier block
+        for (final TieredBlock tieredBlock : tieredBlocks) {
+
+            tierNum++;
+            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;
+                targetTierNum = tierNum;
+                break;
+            }
+        }
+        final int lastBlockTierSize = targetBlock.getSize().intValue();
+        final int nbBlocks = units.intValue() / lastBlockTierSize + (units.intValue() % lastBlockTierSize == 0 ? 0 : 1);
+
+        return new UsageConsumableInArrearTierUnitDetail(targetTierNum, targetBlock.getUnit().getName(), targetBlock.getPrice().getPrice(getCurrency()), targetBlock.getSize().intValue(), nbBlocks);
+    }
+
+    @Override
+    public String toString() {
+        final StringBuilder sb = new StringBuilder("ContiguousIntervalConsumableUsageInArrear{");
+        sb.append("transitionTimes=").append(transitionTimes);
+        sb.append(", billingEvents=").append(billingEvents);
+        sb.append(", rawSubscriptionUsage=").append(rawSubscriptionUsage);
+        sb.append(", rawUsageStartDate=").append(rawUsageStartDate);
+        sb.append('}');
+        return sb.toString();
+    }
+
+    public static List<UsageConsumableInArrearTierUnitDetail> fromJson(String itemDetails) {
+        List<UsageConsumableInArrearTierUnitDetail> toBeBilledUsageConsumableInArrearTierUnitDetails = null;
+        if (itemDetails != null) {
+            try {
+                toBeBilledUsageConsumableInArrearTierUnitDetails = objectMapper.readValue(itemDetails, new TypeReference<List<UsageConsumableInArrearTierUnitDetail>>() {});
+            } catch (IOException e) {
+                Preconditions.checkState(false, e.getMessage());
+            }
+        }
+
+        return toBeBilledUsageConsumableInArrearTierUnitDetails;
+    }
+
+}
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 2bf2213..077db78 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
@@ -50,7 +50,6 @@ import org.killbill.billing.invoice.usage.details.UsageCapacityInArrearDetail;
 import org.killbill.billing.invoice.usage.details.UsageConsumableInArrearDetail;
 import org.killbill.billing.invoice.usage.details.UsageConsumableInArrearTierUnitDetail;
 import org.killbill.billing.invoice.usage.details.UsageInArrearDetail;
-import org.killbill.billing.invoice.usage.details.UsageInArrearDetailInitializer;
 import org.killbill.billing.invoice.usage.details.UsageInArrearTierUnitDetail;
 import org.killbill.billing.junction.BillingEvent;
 import org.killbill.billing.usage.RawUsage;
@@ -67,7 +66,6 @@ import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 
@@ -80,24 +78,24 @@ import static org.killbill.billing.invoice.usage.UsageUtils.getConsumableInArrea
  * There is one such class per subscriptionId, matching a given in arrear/consumable usage section and
  * referenced through a contiguous list of billing events.
  */
-public class ContiguousIntervalUsageInArrear {
+public abstract class ContiguousIntervalUsageInArrear {
 
     private static final Logger log = LoggerFactory.getLogger(ContiguousIntervalUsageInArrear.class);
 
-    private final List<LocalDate> transitionTimes;
-    private final List<BillingEvent> billingEvents;
-
-    private final Usage usage;
-    private final Set<String> unitTypes;
-    private final List<RawUsage> rawSubscriptionUsage;
-    private final LocalDate targetDate;
-    private final UUID accountId;
-    private final UUID invoiceId;
-    private final AtomicBoolean isBuilt;
-    private final LocalDate rawUsageStartDate;
-    private final InternalTenantContext internalTenantContext;
-    private final UsageDetailMode usageDetailMode;
-    private static final ObjectMapper objectMapper = new ObjectMapper();
+    protected final List<LocalDate> transitionTimes;
+    protected final List<BillingEvent> billingEvents;
+
+    protected final Usage usage;
+    protected final Set<String> unitTypes;
+    protected final List<RawUsage> rawSubscriptionUsage;
+    protected final LocalDate targetDate;
+    protected final UUID accountId;
+    protected final UUID invoiceId;
+    protected final AtomicBoolean isBuilt;
+    protected final LocalDate rawUsageStartDate;
+    protected final InternalTenantContext internalTenantContext;
+    protected final UsageDetailMode usageDetailMode;
+    protected static final ObjectMapper objectMapper = new ObjectMapper();
 
     public ContiguousIntervalUsageInArrear(final Usage usage,
                                            final UUID accountId,
@@ -212,82 +210,20 @@ public class ContiguousIntervalUsageInArrear {
 
             final List<RolledUpUnit> rolledUpUnits = ru.getRolledUpUnits();
 
-            UsageInArrearDetail toBeBilledUsageDetails = new UsageInArrearDetailInitializer();
-            BigDecimal toBeBilledUsage;
-            if (usage.getUsageType() == UsageType.CAPACITY) {
-                // In CAPACITY mode we build across units
-                toBeBilledUsageDetails = computeToBeBilledCapacityInArrear(rolledUpUnits);
-            } else /* UsageType.CONSUMABLE */ {
-
-                final Map<String, List<UsageConsumableInArrearTierUnitDetail>> previousUnitsUsage;
-                if (areAllBilledItemsWithDetails) {
-                    previousUnitsUsage = new HashMap<String, List<UsageConsumableInArrearTierUnitDetail>>();
-                    for (RolledUpUnit cur : rolledUpUnits) {
-                        final List<UsageConsumableInArrearTierUnitDetail> usageInArrearDetailForUnitType = getUsageConsumableInArrearDetailForUnitType(billedItems, cur.getUnitType());
-                        previousUnitsUsage.put(cur.getUnitType(), usageInArrearDetailForUnitType);
-                    }
-                } else {
-                    previousUnitsUsage = ImmutableMap.of();
-                }
+            final UsageInArrearDetail toBeBilledUsageDetails = getToBeBilledUsageDetails(rolledUpUnits, billedItems, areAllBilledItemsWithDetails);
+            final BigDecimal toBeBilledUsage = toBeBilledUsageDetails.getAmount();
 
-                // In CONSUMABLE mode we build for each unit independently, hence *for* loop
-                final List<UsageConsumableInArrearTierUnitDetail> usageConsumableInArrearTierUnitDetails = new ArrayList<UsageConsumableInArrearTierUnitDetail>();
-                for (final RolledUpUnit cur : rolledUpUnits) {
-                    if (!unitTypes.contains(cur.getUnitType())) {
-                        log.warn("ContiguousIntervalConsumableInArrear is skipping unitType " + cur.getUnitType());
-                        continue;
-                    }
-                    final List<UsageConsumableInArrearTierUnitDetail> previousUsage = previousUnitsUsage.containsKey(cur.getUnitType()) ? previousUnitsUsage.get(cur.getUnitType()) : ImmutableList.<UsageConsumableInArrearTierUnitDetail>of();
-
-                    usageConsumableInArrearTierUnitDetails.addAll(computeToBeBilledConsumableInArrear(cur, previousUsage, areAllBilledItemsWithDetails));
-                }
-                toBeBilledUsageDetails = new UsageConsumableInArrearDetail(usageConsumableInArrearTierUnitDetails);
-            }
-            // Based on details computed, compute total (proposed) amount
-            toBeBilledUsage = toBeBilledUsageDetails.getAmount();
-
-            // Compute final amount by subtracting  amount that was already billed.
-            if (!billedItems.iterator().hasNext() || billedUsage.compareTo(toBeBilledUsage) < 0) {
-                final BigDecimal amountToBill;
-                if (areAllBilledItemsWithDetails && usage.getUsageType() == UsageType.CONSUMABLE) {
-                    amountToBill = toBeBilledUsage;
-                } else {
-                    amountToBill = toBeBilledUsage.subtract(billedUsage);
-                }
+            populateResults(ru.getStart(), ru.getEnd(), billedItems, billedUsage, toBeBilledUsage, toBeBilledUsageDetails, areAllBilledItemsWithDetails, result);
 
-                if (amountToBill.compareTo(BigDecimal.ZERO) > 0) {
-                    if (UsageDetailMode.DETAIL == usageDetailMode && usage.getUsageType() == UsageType.CONSUMABLE) {
-                        for (UsageConsumableInArrearTierUnitDetail toBeBilledUsageDetail : ((UsageConsumableInArrearDetail) toBeBilledUsageDetails).getTierDetails()) {
-                            final InvoiceItem item = new UsageInvoiceItem(invoiceId, accountId, getBundleId(), getSubscriptionId(), getPlanName(),
-                                                                          getPhaseName(), usage.getName(), ru.getStart(), ru.getEnd(), toBeBilledUsageDetail.getAmount(), toBeBilledUsageDetail.getTierPrice(), getCurrency(), toBeBilledUsageDetail.getQuantity(), null);
-                            result.add(item);
-                        }
-                    } else {
-                        final String itemDetails = areAllBilledItemsWithDetails ? toBeBilledUsageDetails.toJson(objectMapper) : null;
-                        final InvoiceItem item = new UsageInvoiceItem(invoiceId, accountId, getBundleId(), getSubscriptionId(), getPlanName(),
-                                                                      getPhaseName(), usage.getName(), ru.getStart(), ru.getEnd(), amountToBill, null, getCurrency(), null, itemDetails);
-                        result.add(item);
-                    }
-                }
-            }
         }
         final LocalDate nextNotificationdate = computeNextNotificationDate();
         return new UsageInArrearItemsAndNextNotificationDate(result, nextNotificationdate);
     }
 
-    private List<UsageConsumableInArrearTierUnitDetail> getUsageConsumableInArrearDetailForUnitType(final Iterable<InvoiceItem> billedItems, final String unitType) {
+    protected abstract void populateResults(final LocalDate startDate, final LocalDate endDate, final Iterable<InvoiceItem> billedItems, final BigDecimal billedUsage, final BigDecimal toBeBilledUsage, final UsageInArrearDetail toBeBilledUsageDetails, final boolean areAllBilledItemsWithDetails, final List<InvoiceItem> result);
 
-        final List<UsageConsumableInArrearTierUnitDetail> result = new ArrayList<UsageConsumableInArrearTierUnitDetail>();
-        for (final InvoiceItem bi : billedItems) {
-            final List<UsageConsumableInArrearTierUnitDetail> billedUsageItemDetails = fromJson(bi.getItemDetails());
-            for (final UsageConsumableInArrearTierUnitDetail curDetail : billedUsageItemDetails) {
-                if (curDetail.getTierUnit().equals(unitType)) {
-                    result.add(curDetail);
-                }
-            }
-        }
-        return result;
-    }
+
+    protected abstract UsageInArrearDetail getToBeBilledUsageDetails(final List<RolledUpUnit> rolledUpUnits, final Iterable<InvoiceItem> billedItems, final boolean areAllBilledItemsWithDetails) throws CatalogApiException;
 
     private boolean areAllBilledItemsWithDetails(final Iterable<InvoiceItem> billedItems) {
         boolean atLeastOneItemWithoutDetails = Iterables.any(billedItems, new Predicate<InvoiceItem>() {
@@ -475,97 +411,6 @@ public class ContiguousIntervalUsageInArrear {
         return null;
     }
 
-    /**
-     * @param roUnit the rolled up unit for the period
-     * @return the price amount that should be billed for that period/unitType
-     * @throws CatalogApiException
-     */
-    @VisibleForTesting
-    List<UsageConsumableInArrearTierUnitDetail> computeToBeBilledConsumableInArrear(final RolledUpUnit roUnit, final List<UsageConsumableInArrearTierUnitDetail> previousUsage, final boolean areAllBilledItemsWithDetails) throws CatalogApiException {
-
-        Preconditions.checkState(isBuilt.get());
-        final List<TieredBlock> tieredBlocks = getConsumableInArrearTieredBlocks(usage, roUnit.getUnitType());
-
-        switch (usage.getTierBlockPolicy()) {
-            case ALL_TIERS:
-                return computeToBeBilledConsumableInArrearWith_ALL_TIERS(tieredBlocks, previousUsage, roUnit.getAmount());
-            case TOP_TIER:
-                return Arrays.asList(computeToBeBilledConsumableInArrearWith_TOP_TIER(tieredBlocks, previousUsage, roUnit.getAmount(), areAllBilledItemsWithDetails));
-            default:
-                throw new IllegalStateException("Unknown TierBlockPolicy " + usage.getTierBlockPolicy());
-        }
-    }
-
-    List<UsageConsumableInArrearTierUnitDetail> computeToBeBilledConsumableInArrearWith_ALL_TIERS(final List<TieredBlock> tieredBlocks, final List<UsageConsumableInArrearTierUnitDetail> previousUsage, final Long units) throws CatalogApiException {
-
-        List<UsageConsumableInArrearTierUnitDetail> toBeBilledDetails = Lists.newLinkedList();
-        int remainingUnits = units.intValue();
-        int tierNum = 0;
-
-        final int lastPreviousUsageTier = previousUsage.size(); // we count tier from 1, 2, ...
-        final boolean hasPreviousUsage = lastPreviousUsageTier > 0;
-
-        for (final TieredBlock tieredBlock : tieredBlocks) {
-
-            tierNum++;
-            final int blockTierSize = tieredBlock.getSize().intValue();
-            final int tmp = remainingUnits / blockTierSize + (remainingUnits % blockTierSize == 0 ? 0 : 1);
-            int nbUsedTierBlocks;
-            if (tmp > tieredBlock.getMax()) {
-                nbUsedTierBlocks = tieredBlock.getMax().intValue();
-                remainingUnits -= tieredBlock.getMax() * blockTierSize;
-            } else {
-                nbUsedTierBlocks = tmp;
-                remainingUnits = 0;
-            }
-
-            if (nbUsedTierBlocks > 0) {
-                if (hasPreviousUsage) {
-                    final Integer previousUsageQuantity = tierNum <= lastPreviousUsageTier ? previousUsage.get(tierNum - 1).getQuantity() : 0;
-                    if (tierNum < lastPreviousUsageTier) {
-                        Preconditions.checkState(nbUsedTierBlocks == previousUsageQuantity, String.format("Expected usage for tier='%d', unit='%s' to be full, instead found units='[%d/%d]'",
-                                                                                                          tierNum, tieredBlock.getUnit().getName(), nbUsedTierBlocks, previousUsageQuantity));
-                    } else {
-                        Preconditions.checkState(nbUsedTierBlocks - previousUsageQuantity >= 0, String.format("Expected usage for tier='%d', unit='%s' to contain at least as mush as current usage, instead found units='[%d/%d]",
-                                                                                                              tierNum, tieredBlock.getUnit().getName(), nbUsedTierBlocks, previousUsageQuantityq));
-                    }
-                    nbUsedTierBlocks = nbUsedTierBlocks - previousUsageQuantity;
-                }
-                if (nbUsedTierBlocks > 0) {
-                    toBeBilledDetails.add(new UsageConsumableInArrearTierUnitDetail(tierNum, tieredBlock.getUnit().getName(), tieredBlock.getPrice().getPrice(getCurrency()), blockTierSize, nbUsedTierBlocks));
-                }
-            }
-        }
-        return toBeBilledDetails;
-    }
-
-    UsageConsumableInArrearTierUnitDetail computeToBeBilledConsumableInArrearWith_TOP_TIER(final List<TieredBlock> tieredBlocks, final List<UsageConsumableInArrearTierUnitDetail> previousUsage, final Long units, final boolean areAllBilledItemsWithDetails) throws CatalogApiException {
-
-        int remainingUnits = units.intValue();
-
-        // By default last last tierBlock
-        TieredBlock targetBlock = tieredBlocks.get(tieredBlocks.size() - 1);
-        int targetTierNum = tieredBlocks.size();
-        int tierNum = 0;
-        // Loop through all tier block
-        for (final TieredBlock tieredBlock : tieredBlocks) {
-
-            tierNum++;
-            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;
-                targetTierNum = tierNum;
-                break;
-            }
-        }
-        final int lastBlockTierSize = targetBlock.getSize().intValue();
-        final int nbBlocks = units.intValue() / lastBlockTierSize + (units.intValue() % lastBlockTierSize == 0 ? 0 : 1);
-
-        return new UsageConsumableInArrearTierUnitDetail(targetTierNum, targetBlock.getUnit().getName(), targetBlock.getPrice().getPrice(getCurrency()), targetBlock.getSize().intValue(), nbBlocks);
-    }
 
     /**
      * @param filteredUsageForInterval the list of invoiceItem to consider
@@ -695,5 +540,4 @@ public class ContiguousIntervalUsageInArrear {
 
         return toBeBilledUsageConsumableInArrearTierUnitDetails;
     }
-
 }
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/usage/details/UsageCapacityInArrearDetail.java b/invoice/src/main/java/org/killbill/billing/invoice/usage/details/UsageCapacityInArrearDetail.java
index 2030690..e80b7a7 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/usage/details/UsageCapacityInArrearDetail.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/usage/details/UsageCapacityInArrearDetail.java
@@ -51,11 +51,10 @@ public class UsageCapacityInArrearDetail implements UsageInArrearDetail {
         return tierDetails;
     }
 
-    // TODO STEPH optimize
     @Override
     public String toJson(final ObjectMapper objectMapper) {
         String result = null;
-        if (tierDetails != null && tierDetails.size() > 0){
+        if (tierDetails != null && tierDetails.size() > 0) {
             try {
                 result = objectMapper.writeValueAsString(tierDetails);
             } catch (JsonProcessingException e) {
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/usage/SubscriptionUsageInArrear.java b/invoice/src/main/java/org/killbill/billing/invoice/usage/SubscriptionUsageInArrear.java
index b104127..f28d4c9 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/usage/SubscriptionUsageInArrear.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/usage/SubscriptionUsageInArrear.java
@@ -154,7 +154,10 @@ public class SubscriptionUsageInArrear {
                 // Add inflight usage interval if non existent
                 ContiguousIntervalUsageInArrear existingInterval = inFlightInArrearUsageIntervals.get(usage.getName());
                 if (existingInterval == null) {
-                    existingInterval = new ContiguousIntervalUsageInArrear(usage, accountId, invoiceId, rawSubscriptionUsage, targetDate, rawUsageStartDate, usageDetailMode, internalTenantContext);
+                    existingInterval = usage.getUsageType() == UsageType.CAPACITY ?
+                                       new ContiguousIntervalCapacityUsageInArrear(usage, accountId, invoiceId, rawSubscriptionUsage, targetDate, rawUsageStartDate, usageDetailMode, internalTenantContext) :
+                                       new ContiguousIntervalConsumableUsageInArrear(usage, accountId, invoiceId, rawSubscriptionUsage, targetDate, rawUsageStartDate, usageDetailMode, internalTenantContext);
+
                     inFlightInArrearUsageIntervals.put(usage.getName(), existingInterval);
                 }
                 // Add billing event for that usage interval
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalCapacityInArrear.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalCapacityInArrear.java
index 2507d7d..c365295 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalCapacityInArrear.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalCapacityInArrear.java
@@ -86,7 +86,7 @@ public class TestContiguousIntervalCapacityInArrear extends TestUsageInArrearBas
 
         final LocalDate targetDate = startDate.plusDays(1);
 
-        final ContiguousIntervalUsageInArrear intervalCapacityInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+        final ContiguousIntervalUsageInArrear intervalCapacityInArrear = createContiguousIntervalCapacityInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
                                                                                                                                              BillingPeriod.MONTHLY,
                                                                                                                                              Collections.<Usage>emptyList())
@@ -144,7 +144,7 @@ public class TestContiguousIntervalCapacityInArrear extends TestUsageInArrearBas
 
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
 
-        final ContiguousIntervalUsageInArrear intervalCapacityInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+        final ContiguousIntervalUsageInArrear intervalCapacityInArrear = createContiguousIntervalCapacityInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
                                                                                                                                                   BillingPeriod.MONTHLY,
                                                                                                                                                   Collections.<Usage>emptyList())
@@ -219,7 +219,7 @@ public class TestContiguousIntervalCapacityInArrear extends TestUsageInArrearBas
         final BillingEvent event1 = createMockBillingEvent(startDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),BillingPeriod.MONTHLY, Collections.<Usage>emptyList());
         final BillingEvent event2 = createMockBillingEvent(endDate.toDateTimeAtStartOfDay(DateTimeZone.UTC), BillingPeriod.MONTHLY, Collections.<Usage>emptyList());
 
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, rawUsages, targetDate, true, event1, event2);
+        final ContiguousIntervalCapacityUsageInArrear intervalCapacityInArrear = createContiguousIntervalCapacityInArrear(usage, rawUsages, targetDate, true, event1, event2);
 
         final List<InvoiceItem> invoiceItems = new ArrayList<InvoiceItem>();
         final InvoiceItem ii1 = new UsageInvoiceItem(invoiceId, accountId, bundleId, subscriptionId, planName, phaseName, usage.getName(), startDate, firstBCDDate, BigDecimal.ONE, currency);
@@ -228,7 +228,7 @@ public class TestContiguousIntervalCapacityInArrear extends TestUsageInArrearBas
         final InvoiceItem ii2 = new UsageInvoiceItem(invoiceId, accountId, bundleId, subscriptionId, planName, phaseName, usage.getName(), firstBCDDate, endDate, BigDecimal.ONE, currency);
         invoiceItems.add(ii2);
 
-        final UsageInArrearItemsAndNextNotificationDate usageResult = intervalConsumableInArrear.computeMissingItemsAndNextNotificationDate(invoiceItems);
+        final UsageInArrearItemsAndNextNotificationDate usageResult = intervalCapacityInArrear.computeMissingItemsAndNextNotificationDate(invoiceItems);
         final List<InvoiceItem> rawResults = usageResult.getInvoiceItems();
         assertEquals(rawResults.size(), 4);
 
@@ -408,9 +408,9 @@ public class TestContiguousIntervalCapacityInArrear extends TestUsageInArrearBas
         final BillingEvent event1 = createMockBillingEvent(startDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),BillingPeriod.MONTHLY, Collections.<Usage>emptyList());
         final BillingEvent event2 = createMockBillingEvent(endDate.toDateTimeAtStartOfDay(DateTimeZone.UTC), BillingPeriod.MONTHLY, Collections.<Usage>emptyList());
 
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, rawUsages, targetDate, true, usageDetailMode, event1, event2);
+        final ContiguousIntervalCapacityUsageInArrear intervalCapacityInArrear = createContiguousIntervalCapacityInArrear(usage, rawUsages, targetDate, true, usageDetailMode, event1, event2);
 
-        final UsageInArrearItemsAndNextNotificationDate usageResult = intervalConsumableInArrear.computeMissingItemsAndNextNotificationDate(existingItems);
+        final UsageInArrearItemsAndNextNotificationDate usageResult = intervalCapacityInArrear.computeMissingItemsAndNextNotificationDate(existingItems);
         final List<InvoiceItem> rawResults = usageResult.getInvoiceItems();
         final List<InvoiceItem> result = ImmutableList.copyOf(Iterables.filter(rawResults, new Predicate<InvoiceItem>() {
             @Override
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 70909e6..87de942 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
@@ -140,7 +140,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
 
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+        final ContiguousIntervalConsumableUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
                                                                                                                                              BillingPeriod.MONTHLY,
                                                                                                                                              Collections.<Usage>emptyList())
@@ -168,7 +168,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
 
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+        final ContiguousIntervalConsumableUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
                                                                                                                                              BillingPeriod.MONTHLY,
                                                                                                                                              Collections.<Usage>emptyList())
@@ -200,7 +200,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
 
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+        final ContiguousIntervalConsumableUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
                                                                                                                                              BillingPeriod.MONTHLY,
                                                                                                                                              Collections.<Usage>emptyList())
@@ -247,7 +247,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
 
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+        final ContiguousIntervalConsumableUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
                                                                                                                                              BillingPeriod.MONTHLY,
                                                                                                                                              Collections.<Usage>emptyList())
@@ -426,7 +426,11 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
         final BillingEvent event2 = createMockBillingEvent(new LocalDate(2014, 10, 16).toDateTimeAtStartOfDay(DateTimeZone.UTC), BillingPeriod.MONTHLY, Collections.<Usage>emptyList());
 
 
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = new ContiguousIntervalUsageInArrear(usage, accountId, invoiceId, rawUsages, targetDate, rawUsageStartDate, usageDetailMode, internalCallContext);
+
+        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = usage.getUsageType() == UsageType.CAPACITY ?
+                                                                           new ContiguousIntervalCapacityUsageInArrear(usage, accountId, invoiceId, rawUsages, targetDate, rawUsageStartDate, usageDetailMode, internalCallContext) :
+                                                                           new ContiguousIntervalConsumableUsageInArrear(usage, accountId, invoiceId, rawUsages, targetDate, rawUsageStartDate, usageDetailMode, internalCallContext);
+
         intervalConsumableInArrear.addBillingEvent(event1);
         intervalConsumableInArrear.addBillingEvent(event2);
 
@@ -443,7 +447,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final DefaultUsage usage = createConsumableInArrearUsage(usageName, BillingPeriod.MONTHLY, TierBlockPolicy.ALL_TIERS, tier);
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
+        final ContiguousIntervalConsumableUsageInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
                                                                                                                       createMockBillingEvent(targetDate.toDateTimeAtStartOfDay(DateTimeZone.UTC),
                                                                                                                                              BillingPeriod.MONTHLY,
                                                                                                                                              Collections.<Usage>emptyList())
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 7a205c1..42c1709 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
@@ -78,12 +78,27 @@ public abstract class TestUsageInArrearBase extends InvoiceTestSuiteNoDB {
         objectMapper = new ObjectMapper();
     }
 
-    protected ContiguousIntervalUsageInArrear createContiguousIntervalConsumableInArrear(final DefaultUsage usage, final List<RawUsage> rawUsages, final LocalDate targetDate, final boolean closedInterval, final BillingEvent... events) {
+
+    protected ContiguousIntervalCapacityUsageInArrear createContiguousIntervalCapacityInArrear(final DefaultUsage usage, final List<RawUsage> rawUsages, final LocalDate targetDate, final boolean closedInterval, final BillingEvent... events) {
+        return createContiguousIntervalCapacityInArrear(usage, rawUsages, targetDate, closedInterval, usageDetailMode, events);
+    }
+
+    protected ContiguousIntervalCapacityUsageInArrear createContiguousIntervalCapacityInArrear(final DefaultUsage usage, final List<RawUsage> rawUsages, final LocalDate targetDate, final boolean closedInterval, UsageDetailMode detailMode, final BillingEvent... events) {
+        final ContiguousIntervalCapacityUsageInArrear intervalCapacityInArrear = new ContiguousIntervalCapacityUsageInArrear(usage, accountId, invoiceId, rawUsages, targetDate, new LocalDate(events[0].getEffectiveDate()),  detailMode, internalCallContext);
+        for (final BillingEvent event : events) {
+            intervalCapacityInArrear.addBillingEvent(event);
+        }
+        intervalCapacityInArrear.build(closedInterval);
+        return intervalCapacityInArrear;
+    }
+
+
+    protected ContiguousIntervalConsumableUsageInArrear createContiguousIntervalConsumableInArrear(final DefaultUsage usage, final List<RawUsage> rawUsages, final LocalDate targetDate, final boolean closedInterval, final BillingEvent... events) {
         return createContiguousIntervalConsumableInArrear(usage, rawUsages, targetDate, closedInterval, usageDetailMode, events);
     }
 
-    protected ContiguousIntervalUsageInArrear createContiguousIntervalConsumableInArrear(final DefaultUsage usage, final List<RawUsage> rawUsages, final LocalDate targetDate, final boolean closedInterval, UsageDetailMode detailMode, final BillingEvent... events) {
-        final ContiguousIntervalUsageInArrear intervalConsumableInArrear = new ContiguousIntervalUsageInArrear(usage, accountId, invoiceId, rawUsages, targetDate, new LocalDate(events[0].getEffectiveDate()), detailMode, internalCallContext);
+    protected ContiguousIntervalConsumableUsageInArrear createContiguousIntervalConsumableInArrear(final DefaultUsage usage, final List<RawUsage> rawUsages, final LocalDate targetDate, final boolean closedInterval, UsageDetailMode detailMode, final BillingEvent... events) {
+        final ContiguousIntervalConsumableUsageInArrear intervalConsumableInArrear = new ContiguousIntervalConsumableUsageInArrear(usage, accountId, invoiceId, rawUsages, targetDate, new LocalDate(events[0].getEffectiveDate()), detailMode, internalCallContext);
         for (final BillingEvent event : events) {
             intervalConsumableInArrear.addBillingEvent(event);
         }