killbill-aplcache

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/overdue/TestOverdueBase.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/overdue/TestOverdueBase.java
index ccc25d5..5f5b4e2 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/overdue/TestOverdueBase.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/overdue/TestOverdueBase.java
@@ -27,11 +27,12 @@ import org.killbill.billing.account.api.Account;
 import org.killbill.billing.beatrix.integration.BeatrixIntegrationModule;
 import org.killbill.billing.beatrix.integration.TestIntegrationBase;
 import org.killbill.billing.catalog.api.BillingPeriod;
+import org.killbill.billing.entitlement.api.BlockingState;
 import org.killbill.billing.entitlement.api.BlockingStateType;
 import org.killbill.billing.entitlement.api.SubscriptionBundle;
 import org.killbill.billing.overdue.OverdueService;
-import org.killbill.billing.overdue.api.DefaultOverdueInternalApi;
 import org.killbill.billing.overdue.config.DefaultOverdueConfig;
+import org.killbill.billing.overdue.wrapper.OverdueWrapper;
 import org.killbill.billing.payment.api.PaymentMethodPlugin;
 import org.killbill.billing.payment.api.TestPaymentMethodPluginBase;
 import org.killbill.xmlloader.XMLLoader;
@@ -81,11 +82,15 @@ public abstract class TestOverdueBase extends TestIntegrationBase {
             await().atMost(10, SECONDS).until(new Callable<Boolean>() {
                 @Override
                 public Boolean call() throws Exception {
-                    return expected.equals(blockingApi.getBlockingStateForService(account.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalCallContext).getStateName());
+                    final BlockingState blockingStateForService = blockingApi.getBlockingStateForService(account.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalCallContext);
+                    final String stateName = blockingStateForService != null ? blockingStateForService.getStateName() : OverdueWrapper.CLEAR_STATE_NAME;
+                    return expected.equals(stateName);
                 }
             });
         } catch (final Exception e) {
-            Assert.assertEquals(blockingApi.getBlockingStateForService(account.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalCallContext).getStateName(), expected, "Got exception: " + e.toString());
+            final BlockingState blockingStateForService = blockingApi.getBlockingStateForService(account.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalCallContext);
+            final String stateName = blockingStateForService != null ? blockingStateForService.getStateName() : OverdueWrapper.CLEAR_STATE_NAME;
+            Assert.assertEquals(stateName, expected, "Got exception: " + e.toString());
         }
     }
 }

diff --git a/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java b/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java
index 4444547..a2701e4 100644
--- a/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java
+++ b/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java
@@ -26,6 +26,7 @@ import org.killbill.billing.ErrorCode;
 import org.killbill.billing.account.api.ImmutableAccountData;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
+import org.killbill.billing.entitlement.api.BlockingState;
 import org.killbill.billing.entitlement.api.BlockingStateType;
 import org.killbill.billing.junction.BlockingInternalApi;
 import org.killbill.billing.overdue.OverdueInternalApi;
@@ -84,7 +85,8 @@ public class DefaultOverdueInternalApi implements OverdueInternalApi {
     public OverdueState getOverdueStateFor(final ImmutableAccountData overdueable, final TenantContext context) throws OverdueException {
         try {
             final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(context);
-            final String stateName = MoreObjects.firstNonNull(accessApi.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalCallContextFactory.createInternalTenantContext(context)).getStateName(), OverdueWrapper.CLEAR_STATE_NAME);
+            final BlockingState blockingStateForService = accessApi.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalCallContextFactory.createInternalTenantContext(context));
+            final String stateName = blockingStateForService != null ? blockingStateForService.getStateName() : OverdueWrapper.CLEAR_STATE_NAME;
             final OverdueConfig overdueConfig = overdueConfigCache.getOverdueConfig(internalTenantContext);
             final OverdueStateSet states = ((DefaultOverdueConfig) overdueConfig).getOverdueStatesAccount();
             return states.findState(stateName);

diff --git a/overdue/src/main/java/org/killbill/billing/overdue/wrapper/OverdueWrapper.java b/overdue/src/main/java/org/killbill/billing/overdue/wrapper/OverdueWrapper.java
index 311db99..95a2b7d 100644
--- a/overdue/src/main/java/org/killbill/billing/overdue/wrapper/OverdueWrapper.java
+++ b/overdue/src/main/java/org/killbill/billing/overdue/wrapper/OverdueWrapper.java
@@ -21,6 +21,7 @@ package org.killbill.billing.overdue.wrapper;
 import org.killbill.billing.account.api.ImmutableAccountData;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
+import org.killbill.billing.entitlement.api.BlockingState;
 import org.killbill.billing.entitlement.api.BlockingStateType;
 import org.killbill.billing.junction.BlockingInternalApi;
 import org.killbill.billing.overdue.OverdueService;
@@ -98,8 +99,8 @@ public class OverdueWrapper {
 
     private OverdueState refreshWithLock(final InternalCallContext context) throws OverdueException, OverdueApiException {
         final BillingState billingState = billingState(context);
-        final String previousOverdueStateName = MoreObjects.firstNonNull(api.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, context).getStateName(), OverdueWrapper.CLEAR_STATE_NAME);
-
+        final BlockingState blockingStateForService = api.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, context);
+        final String previousOverdueStateName = blockingStateForService != null ? blockingStateForService.getStateName() : OverdueWrapper.CLEAR_STATE_NAME;
         final OverdueState currentOverdueState = overdueStateSet.findState(previousOverdueStateName);
         final OverdueState nextOverdueState = overdueStateSet.calculateOverdueState(billingState, clock.getToday(billingState.getAccountTimeZone()));
 
@@ -125,7 +126,8 @@ public class OverdueWrapper {
     }
 
     private void clearWithLock(final InternalCallContext context) throws OverdueException, OverdueApiException {
-        final String previousOverdueStateName = MoreObjects.firstNonNull(api.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, context).getStateName(), OverdueWrapper.CLEAR_STATE_NAME);
+        final BlockingState blockingStateForService = api.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, context);
+        final String previousOverdueStateName = blockingStateForService != null ? blockingStateForService.getStateName() : OverdueWrapper.CLEAR_STATE_NAME;
         final OverdueState previousOverdueState = overdueStateSet.findState(previousOverdueStateName);
         overdueStateApplicator.clear(overdueable, previousOverdueState, overdueStateSet.getClearState(), context);
     }