killbill-aplcache

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
index 2969715..d38cea6 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
@@ -799,18 +799,43 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB {
     }
 
     protected void add_AUTO_PAY_OFF_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
+        add_account_Tag(id, ControlTagType.AUTO_PAY_OFF, type);
+    }
+
+    protected void add_AUTO_INVOICING_OFF_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
+        add_account_Tag(id, ControlTagType.AUTO_INVOICING_OFF, type);
+    }
+
+    protected void add_AUTO_INVOICING_DRAFT_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
+        add_account_Tag(id, ControlTagType.AUTO_INVOICING_DRAFT, type);
+    }
+
+    private void add_account_Tag(final UUID id, final ControlTagType controlTagType, final ObjectType type) throws TagDefinitionApiException, TagApiException {
         busHandler.pushExpectedEvent(NextEvent.TAG);
-        tagUserApi.addTag(id, type, ControlTagType.AUTO_PAY_OFF.getId(), callContext);
+        tagUserApi.addTag(id, type, controlTagType.getId(), callContext);
         assertListenerStatus();
-
         final List<Tag> tags = tagUserApi.getTagsForObject(id, type, false, callContext);
         assertEquals(tags.size(), 1);
     }
 
     protected void remove_AUTO_PAY_OFF_Tag(final UUID id, final ObjectType type, final NextEvent... additionalEvents) throws TagDefinitionApiException, TagApiException {
+        remove_account_Tag(id, ControlTagType.AUTO_PAY_OFF, type, additionalEvents);
+    }
+
+
+    protected void remove_AUTO_INVOICING_OFF_Tag(final UUID id, final ObjectType type, final NextEvent... additionalEvents) throws TagDefinitionApiException, TagApiException {
+        remove_account_Tag(id, ControlTagType.AUTO_INVOICING_OFF, type, additionalEvents);
+    }
+
+    protected void remove_AUTO_INVOICING_DRAFT_Tag(final UUID id, final ObjectType type, final NextEvent... additionalEvents) throws TagDefinitionApiException, TagApiException {
+        remove_account_Tag(id, ControlTagType.AUTO_INVOICING_DRAFT, type, additionalEvents);
+    }
+
+
+    private  void remove_account_Tag(final UUID id, final ControlTagType controlTagType, final ObjectType type, final NextEvent... additionalEvents) throws TagDefinitionApiException, TagApiException {
         busHandler.pushExpectedEvent(NextEvent.TAG);
         busHandler.pushExpectedEvents(additionalEvents);
-        tagUserApi.removeTag(id, type, ControlTagType.AUTO_PAY_OFF.getId(), callContext);
+        tagUserApi.removeTag(id, type, controlTagType.getId(), callContext);
         assertListenerStatus();
     }
 

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceDraft.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceDraft.java
index e223f22..5353676 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceDraft.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceDraft.java
@@ -106,9 +106,7 @@ public class TestIntegrationWithAutoInvoiceDraft extends TestIntegrationBase {
         assertListenerStatus();
 
 
-        busHandler.pushExpectedEvents(NextEvent.TAG);
         remove_AUTO_INVOICING_DRAFT_Tag(account.getId(), ObjectType.ACCOUNT);
-        assertListenerStatus();
 
         busHandler.pushExpectedEvents(NextEvent.INVOICE, NextEvent.PAYMENT, NextEvent.INVOICE_PAYMENT);
         clock.addMonths(1);
@@ -167,9 +165,7 @@ public class TestIntegrationWithAutoInvoiceDraft extends TestIntegrationBase {
         final BigDecimal accountBalance = invoiceApi.getAccountBalance(account.getId(), callContext);
         assertEquals(accountBalance.compareTo(BigDecimal.ZERO), 0);
 
-        busHandler.pushExpectedEvents(NextEvent.TAG);
         remove_AUTO_INVOICING_DRAFT_Tag(account.getId(), ObjectType.ACCOUNT);
-        assertListenerStatus();
 
         busHandler.pushExpectedEvents(NextEvent.INVOICE, NextEvent.PAYMENT, NextEvent.INVOICE_PAYMENT);
         clock.addMonths(1);
@@ -199,21 +195,4 @@ public class TestIntegrationWithAutoInvoiceDraft extends TestIntegrationBase {
 
     }
 
-
-    private void add_AUTO_INVOICING_DRAFT_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
-        add_account_Tag(id, ControlTagType.AUTO_INVOICING_DRAFT, type);
-    }
-
-    private void add_account_Tag(final UUID id, final ControlTagType controlTagType, final ObjectType type) throws TagDefinitionApiException, TagApiException {
-        busHandler.pushExpectedEvent(NextEvent.TAG);
-        tagApi.addTag(id, type, controlTagType.getId(), callContext);
-        assertListenerStatus();
-        final List<Tag> tags = tagApi.getTagsForObject(id, type, false, callContext);
-        assertEquals(tags.size(), 1);
-    }
-
-    private void remove_AUTO_INVOICING_DRAFT_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
-        tagApi.removeTag(id, type, ControlTagType.AUTO_INVOICING_DRAFT.getId(), callContext);
-    }
-
 }

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java
index 5b5d4c2..3cd5c9b 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java
@@ -92,9 +92,7 @@ public class TestIntegrationWithAutoInvoiceOffTag extends TestIntegrationBase {
         invoices = invoiceApi.getInvoicesByAccount(account.getId(), false, callContext);
         assertEquals(invoices.size(), 0);
 
-        busHandler.pushExpectedEvents(NextEvent.TAG, NextEvent.INVOICE, NextEvent.PAYMENT, NextEvent.INVOICE_PAYMENT);
-        remove_AUTO_INVOICING_OFF_Tag(account.getId(), ObjectType.ACCOUNT);
-        assertListenerStatus();
+        remove_AUTO_INVOICING_OFF_Tag(account.getId(), ObjectType.ACCOUNT, NextEvent.INVOICE, NextEvent.PAYMENT, NextEvent.INVOICE_PAYMENT);
 
         invoices = invoiceApi.getInvoicesByAccount(account.getId(), false, callContext);
         assertEquals(invoices.size(), 1);
@@ -145,19 +143,4 @@ public class TestIntegrationWithAutoInvoiceOffTag extends TestIntegrationBase {
         assertEquals(invoices.size(), 3); // Only one additional invoice generated
     }
 
-    private void add_AUTO_INVOICING_OFF_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
-        add_account_Tag(id, ControlTagType.AUTO_INVOICING_OFF, type);
-    }
-
-    private void add_account_Tag(final UUID id, final ControlTagType controlTagType, final ObjectType type) throws TagDefinitionApiException, TagApiException {
-        busHandler.pushExpectedEvent(NextEvent.TAG);
-        tagApi.addTag(id, type, controlTagType.getId(), callContext);
-        assertListenerStatus();
-        final List<Tag> tags = tagApi.getTagsForObject(id, type, false, callContext);
-        assertEquals(tags.size(), 1);
-    }
-
-    private void remove_AUTO_INVOICING_OFF_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
-        tagApi.removeTag(id, type, ControlTagType.AUTO_INVOICING_OFF.getId(), callContext);
-    }
 }