killbill-memoizeit

diff --git a/invoice/src/test/java/com/ning/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java b/invoice/src/test/java/com/ning/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java
index 0f6d556..cdc620b 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java
@@ -102,6 +102,10 @@ public class TestDefaultInvoiceMigrationApi extends InvoicingTestBase {
     @Inject
     private BillingApi billingApi;
 
+    @Inject
+    private AccountUserApi accountUserApi;
+
+    private Account account;
     private UUID accountId;
     private UUID subscriptionId;
     private DateTime date_migrated;
@@ -152,9 +156,14 @@ public class TestDefaultInvoiceMigrationApi extends InvoicingTestBase {
         cleanup();
         busService.getBus().start();
 
+        account = Mockito.mock(Account.class);
+        Mockito.when(accountUserApi.getAccountById(accountId)).thenReturn(account);
+        Mockito.when(account.getCurrency()).thenReturn(Currency.USD);
+        Mockito.when(account.getId()).thenReturn(accountId);
+        Mockito.when(account.isNotifiedForInvoices()).thenReturn(true);
+
         migrationInvoiceId = createAndCheckMigrationInvoice();
         regularInvoiceId = generateRegularInvoice();
-
     }
 
     @AfterSuite(groups = {"slow"})
@@ -188,13 +197,6 @@ public class TestDefaultInvoiceMigrationApi extends InvoicingTestBase {
     }
 
     private UUID generateRegularInvoice() throws Exception {
-        final AccountUserApi accountUserApi = Mockito.mock(AccountUserApi.class);
-        final Account account = Mockito.mock(Account.class);
-        Mockito.when(accountUserApi.getAccountById(accountId)).thenReturn(account);
-        Mockito.when(account.getCurrency()).thenReturn(Currency.USD);
-        Mockito.when(account.getId()).thenReturn(accountId);
-        Mockito.when(account.isNotifiedForInvoices()).thenReturn(true);
-
         final Subscription subscription = Mockito.mock(Subscription.class);
         Mockito.when(subscription.getId()).thenReturn(subscriptionId);
         Mockito.when(subscription.getBundleId()).thenReturn(new UUID(0L, 0L));
@@ -261,10 +263,8 @@ public class TestDefaultInvoiceMigrationApi extends InvoicingTestBase {
         final BigDecimal balanceOfAllInvoices = migrationInvoice.getBalance().add(regularInvoice.getBalance());
 
         final BigDecimal accountBalance = invoiceUserApi.getAccountBalance(accountId);
-        System.out.println("ACCOUNT balance: " + accountBalance + " should equal the Balance Of All Invoices: " + balanceOfAllInvoices);
+        log.info("ACCOUNT balance: " + accountBalance + " should equal the Balance Of All Invoices: " + balanceOfAllInvoices);
         Assert.assertEquals(accountBalance.compareTo(balanceOfAllInvoices), 0);
-
-
     }
 
     private boolean checkContains(final List<Invoice> invoices, final UUID invoiceId) {

diff --git a/payment/src/test/java/com/ning/billing/payment/glue/PaymentTestModuleWithMocks.java b/payment/src/test/java/com/ning/billing/payment/glue/PaymentTestModuleWithMocks.java
index ca497f6..c253ff3 100644
--- a/payment/src/test/java/com/ning/billing/payment/glue/PaymentTestModuleWithMocks.java
+++ b/payment/src/test/java/com/ning/billing/payment/glue/PaymentTestModuleWithMocks.java
@@ -20,22 +20,26 @@ import java.io.IOException;
 import java.net.URL;
 import java.util.Map;
 import java.util.Properties;
+import java.util.UUID;
 
+import org.mockito.Mockito;
 import org.skife.config.ConfigSource;
 import org.skife.config.SimplePropertyConfigSource;
 
+import com.google.common.collect.ImmutableMap;
 import com.ning.billing.config.PaymentConfig;
 import com.ning.billing.mock.glue.MockInvoiceModule;
 import com.ning.billing.mock.glue.MockNotificationQueueModule;
-import com.ning.billing.mock.glue.TestDbiModule;
 import com.ning.billing.payment.dao.MockPaymentDao;
 import com.ning.billing.payment.dao.PaymentDao;
 import com.ning.billing.payment.provider.MockPaymentProviderPluginModule;
+import com.ning.billing.util.api.TagUserApi;
+import com.ning.billing.util.dao.ObjectType;
 import com.ning.billing.util.globallocker.GlobalLocker;
 import com.ning.billing.util.globallocker.MockGlobalLocker;
 import com.ning.billing.util.glue.BusModule;
 import com.ning.billing.util.glue.BusModule.BusType;
-import com.ning.billing.util.glue.TagStoreModule;
+import com.ning.billing.util.tag.Tag;
 
 import static org.testng.Assert.assertNotNull;
 
@@ -89,8 +93,11 @@ public class PaymentTestModuleWithMocks extends PaymentModule {
         install(new BusModule(BusType.MEMORY));
         install(new MockNotificationQueueModule());
         install(new MockInvoiceModule());
-        install(new TestDbiModule());
-        install(new TagStoreModule());
+
+        final TagUserApi tagUserApi = Mockito.mock(TagUserApi.class);
+        bind(TagUserApi.class).toInstance(tagUserApi);
+        Mockito.when(tagUserApi.getTags(Mockito.<UUID>any(), Mockito.<ObjectType>any())).thenReturn(ImmutableMap.<String, Tag>of());
+
         bind(GlobalLocker.class).to(MockGlobalLocker.class).asEagerSingleton();
     }
 }