killbill-memoizeit

diff --git a/account/src/main/java/com/ning/billing/account/dao/AuditedAccountDao.java b/account/src/main/java/com/ning/billing/account/dao/AuditedAccountDao.java
index facd28f..ebb0a63 100644
--- a/account/src/main/java/com/ning/billing/account/dao/AuditedAccountDao.java
+++ b/account/src/main/java/com/ning/billing/account/dao/AuditedAccountDao.java
@@ -41,7 +41,6 @@ import com.ning.billing.util.customfield.CustomField;
 import com.ning.billing.util.customfield.dao.CustomFieldSqlDao;
 import com.ning.billing.util.bus.Bus;
 import com.ning.billing.util.tag.Tag;
-import com.ning.billing.util.tag.dao.TagSqlDao;
 
 public class AuditedAccountDao implements AccountDao {
     private final AccountSqlDao accountSqlDao;
@@ -243,8 +242,7 @@ public class AuditedAccountDao implements AccountDao {
     }
 
     private void setTagsFromWithinTransaction(final Account account, final AccountSqlDao transactionalDao) {
-        TagSqlDao tagDao = transactionalDao.become(TagSqlDao.class);
-        List<Tag> tags = tagDao.load(account.getId().toString(), account.getObjectName());
+        List<Tag> tags = tagDao.loadTagsFromTransaction(transactionalDao, account.getId(), Account.ObjectType);
         account.clearTags();
 
         if (tags != null) {
@@ -254,7 +252,7 @@ public class AuditedAccountDao implements AccountDao {
 
     private void saveTagsFromWithinTransaction(final Account account, final AccountSqlDao transactionalDao,
                                                final CallContext context) {
-        tagDao.saveTags(transactionalDao, account.getId(), account.getObjectName(), account.getTagList(), context);
+        tagDao.saveTagsFromTransaction(transactionalDao, account.getId(), account.getObjectName(), account.getTagList(), context);
     }
 
     private void saveCustomFieldsFromWithinTransaction(final Account account, final AccountSqlDao transactionalDao,

diff --git a/account/src/test/java/com/ning/billing/account/glue/AccountModuleWithMocks.java b/account/src/test/java/com/ning/billing/account/glue/AccountModuleWithMocks.java
index 5292161..74fe321 100644
--- a/account/src/test/java/com/ning/billing/account/glue/AccountModuleWithMocks.java
+++ b/account/src/test/java/com/ning/billing/account/glue/AccountModuleWithMocks.java
@@ -21,7 +21,7 @@ import com.ning.billing.account.dao.MockAccountDao;
 import com.ning.billing.util.clock.MockClockModule;
 import com.ning.billing.util.glue.CallContextModule;
 import com.ning.billing.util.glue.FieldStoreModule;
-import com.ning.billing.util.glue.TagStoreModule;
+import com.ning.billing.util.tag.MockTagStoreModuleMemory;
 
 public class AccountModuleWithMocks extends AccountModule {
     @Override
@@ -35,7 +35,7 @@ public class AccountModuleWithMocks extends AccountModule {
         super.configure();
         install(new MockClockModule());
         install(new CallContextModule());
-        install(new TagStoreModule());
+        install(new MockTagStoreModuleMemory());
         install(new FieldStoreModule());
     }
 }

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/glue/MockEngineModuleMemory.java b/entitlement/src/test/java/com/ning/billing/entitlement/glue/MockEngineModuleMemory.java
index 786f1e3..8b78c2d 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/glue/MockEngineModuleMemory.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/glue/MockEngineModuleMemory.java
@@ -29,10 +29,10 @@ public class MockEngineModuleMemory extends MockEngineModule {
         bind(EntitlementDao.class).to(MockEntitlementDaoMemory.class).asEagerSingleton();
     }
 
-
     private void installNotificationQueue() {
         bind(NotificationQueueService.class).to(MockNotificationQueueService.class).asEagerSingleton();
     }
+
     @Override
     protected void configure() {
         super.configure();

diff --git a/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java b/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java
index 3c32f4e..c8ddda3 100644
--- a/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java
+++ b/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java
@@ -29,11 +29,11 @@ import com.ning.billing.util.customfield.dao.CustomFieldSqlDao;
 import com.ning.billing.util.tag.ControlTagType;
 import com.ning.billing.util.tag.Tag;
 import com.ning.billing.util.tag.dao.TagDao;
-import com.ning.billing.util.tag.dao.TagSqlDao;
 import org.joda.time.DateTime;
 import org.skife.jdbi.v2.IDBI;
 import org.skife.jdbi.v2.Transaction;
 import org.skife.jdbi.v2.TransactionStatus;
+import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -304,10 +304,8 @@ public class DefaultInvoiceDao implements InvoiceDao {
         }
     }
 
-    private void getTagsWithinTransaction(final Invoice invoice, final InvoiceSqlDao invoiceSqlDao) {
-        TagSqlDao tagSqlDao = invoiceSqlDao.become(TagSqlDao.class);
-        String invoiceId = invoice.getId().toString();
-        List<Tag> tags = tagSqlDao.load(invoiceId, Invoice.ObjectType);
+    private void getTagsWithinTransaction(final Invoice invoice, final Transmogrifier dao) {
+        List<Tag> tags = tagDao.loadTagsFromTransaction(dao, invoice.getId(), Invoice.ObjectType);
         invoice.addTags(tags);
     }
 

diff --git a/util/src/main/java/com/ning/billing/util/customfield/dao/CustomFieldDao.java b/util/src/main/java/com/ning/billing/util/customfield/dao/CustomFieldDao.java
index 9e62a80..a987c4d 100644
--- a/util/src/main/java/com/ning/billing/util/customfield/dao/CustomFieldDao.java
+++ b/util/src/main/java/com/ning/billing/util/customfield/dao/CustomFieldDao.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2010-2011 Ning, Inc.
+ *
+ * Ning 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 com.ning.billing.util.customfield.dao;
 
 import com.ning.billing.util.callcontext.CallContext;

diff --git a/util/src/main/java/com/ning/billing/util/entity/EntityCollectionBase.java b/util/src/main/java/com/ning/billing/util/entity/EntityCollectionBase.java
index 369f255..a240000 100644
--- a/util/src/main/java/com/ning/billing/util/entity/EntityCollectionBase.java
+++ b/util/src/main/java/com/ning/billing/util/entity/EntityCollectionBase.java
@@ -54,7 +54,7 @@ public abstract class EntityCollectionBase<T extends Entity> implements EntityCo
 
     @Override
     public void remove(T entity) {
-        entities.remove(entity);
+        entities.remove(getEntityKey(entity));
     }
 
     @Override

diff --git a/util/src/main/java/com/ning/billing/util/glue/TagDescriptionDaoProvider.java b/util/src/main/java/com/ning/billing/util/glue/TagDescriptionDaoProvider.java
index 560a09f..c05ca29 100644
--- a/util/src/main/java/com/ning/billing/util/glue/TagDescriptionDaoProvider.java
+++ b/util/src/main/java/com/ning/billing/util/glue/TagDescriptionDaoProvider.java
@@ -1,39 +1,39 @@
-/*
- * Copyright 2010-2011 Ning, Inc.
- *
- * Ning 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 com.ning.billing.util.glue;
-
-import com.google.inject.Inject;
-import com.google.inject.Provider;
-import com.ning.billing.util.tag.dao.TagDefinitionSqlDao;
-import org.skife.jdbi.v2.IDBI;
-
-public class TagDescriptionDaoProvider implements Provider<TagDefinitionSqlDao>
-{
-    private final IDBI dbi;
-
-    @Inject
-    public TagDescriptionDaoProvider(final IDBI dbi)
-    {
-        this.dbi = dbi;
-    }
-
-    @Override
-    public TagDefinitionSqlDao get()
-    {
-        return dbi.onDemand(TagDefinitionSqlDao.class);
-    }
-}
+///*
+// * Copyright 2010-2011 Ning, Inc.
+// *
+// * Ning 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 com.ning.billing.util.glue;
+//
+//import com.google.inject.Inject;
+//import com.google.inject.Provider;
+//import com.ning.billing.util.tag.dao.TagDefinitionSqlDao;
+//import org.skife.jdbi.v2.IDBI;
+//
+//public class TagDescriptionDaoProvider implements Provider<TagDefinitionSqlDao>
+//{
+//    private final IDBI dbi;
+//
+//    @Inject
+//    public TagDescriptionDaoProvider(final IDBI dbi)
+//    {
+//        this.dbi = dbi;
+//    }
+//
+//    @Override
+//    public TagDefinitionSqlDao get()
+//    {
+//        return dbi.onDemand(TagDefinitionSqlDao.class);
+//    }
+//}

diff --git a/util/src/main/java/com/ning/billing/util/glue/TagStoreDaoProvider.java b/util/src/main/java/com/ning/billing/util/glue/TagStoreDaoProvider.java
index 7b73f06..1f97f7e 100644
--- a/util/src/main/java/com/ning/billing/util/glue/TagStoreDaoProvider.java
+++ b/util/src/main/java/com/ning/billing/util/glue/TagStoreDaoProvider.java
@@ -1,39 +1,39 @@
-/*
- * Copyright 2010-2011 Ning, Inc.
- *
- * Ning 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 com.ning.billing.util.glue;
-
-import com.google.inject.Inject;
-import com.google.inject.Provider;
-import com.ning.billing.util.tag.dao.TagSqlDao;
-import org.skife.jdbi.v2.IDBI;
-
-public class TagStoreDaoProvider implements Provider<TagSqlDao>
-{
-    private final IDBI dbi;
-
-    @Inject
-    public TagStoreDaoProvider(final IDBI dbi)
-    {
-        this.dbi = dbi;
-    }
-
-    @Override
-    public TagSqlDao get()
-    {
-        return dbi.onDemand(TagSqlDao.class);
-    }
-}
+///*
+// * Copyright 2010-2011 Ning, Inc.
+// *
+// * Ning 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 com.ning.billing.util.glue;
+//
+//import com.google.inject.Inject;
+//import com.google.inject.Provider;
+//import com.ning.billing.util.tag.dao.TagSqlDao;
+//import org.skife.jdbi.v2.IDBI;
+//
+//public class TagStoreDaoProvider implements Provider<TagSqlDao>
+//{
+//    private final IDBI dbi;
+//
+//    @Inject
+//    public TagStoreDaoProvider(final IDBI dbi)
+//    {
+//        this.dbi = dbi;
+//    }
+//
+//    @Override
+//    public TagSqlDao get()
+//    {
+//        return dbi.onDemand(TagSqlDao.class);
+//    }
+//}

diff --git a/util/src/main/java/com/ning/billing/util/glue/TagStoreModule.java b/util/src/main/java/com/ning/billing/util/glue/TagStoreModule.java
index 42cdeff..a1c746e 100644
--- a/util/src/main/java/com/ning/billing/util/glue/TagStoreModule.java
+++ b/util/src/main/java/com/ning/billing/util/glue/TagStoreModule.java
@@ -23,13 +23,9 @@ import com.ning.billing.util.tag.dao.AuditedTagDao;
 import com.ning.billing.util.tag.dao.DefaultTagDefinitionDao;
 import com.ning.billing.util.tag.dao.TagDao;
 import com.ning.billing.util.tag.dao.TagDefinitionDao;
-import com.ning.billing.util.tag.dao.TagDefinitionSqlDao;
-import com.ning.billing.util.tag.dao.TagSqlDao;
-
 public class TagStoreModule extends AbstractModule
 {
     protected void installDaos() {
-        bind(TagDefinitionSqlDao.class).toProvider(TagDescriptionDaoProvider.class).asEagerSingleton();
         bind(TagDefinitionDao.class).to(DefaultTagDefinitionDao.class).asEagerSingleton();
         bind(TagDao.class).to(AuditedTagDao.class).asEagerSingleton();
     }

diff --git a/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java b/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java
index 8cf19cb..6e361f9 100644
--- a/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java
+++ b/util/src/main/java/com/ning/billing/util/tag/dao/AuditedTagDao.java
@@ -32,15 +32,47 @@ import java.util.UUID;
 
 public class AuditedTagDao implements TagDao {
     private final TagSqlDao tagSqlDao;
+    private final TagAuditSqlDao tagAuditSqlDao;
 
     @Inject
     public AuditedTagDao(final IDBI dbi) {
         this.tagSqlDao = dbi.onDemand(TagSqlDao.class);
+        this.tagAuditSqlDao = dbi.onDemand(TagAuditSqlDao.class);
     }
 
     @Override
-    public void saveTags(final Transmogrifier dao, final UUID objectId, final String objectType,
+    public void saveTags(final UUID objectId, final String objectType,
                          final List<Tag> tags, final CallContext context) {
+        // get list of existing tags
+        List<Tag> existingTags = tagSqlDao.load(objectId.toString(), objectType);
+
+        // sort into tags to update (tagsToUpdate), tags to add (tags), and tags to delete (existingTags)
+        Iterator<Tag> tagIterator = tags.iterator();
+        while (tagIterator.hasNext()) {
+            Tag tag = tagIterator.next();
+
+            Iterator<Tag> existingTagIterator = existingTags.iterator();
+            while (existingTagIterator.hasNext()) {
+                Tag existingTag = existingTagIterator.next();
+                if (tag.getTagDefinitionName().equals(existingTag.getTagDefinitionName())) {
+                    // if the tags match, remove from both lists
+                    // in the case of tag, this just means the tag remains associated
+                    tagIterator.remove();
+                    existingTagIterator.remove();
+                }
+            }
+        }
+
+        tagSqlDao.batchInsertFromTransaction(objectId.toString(), objectType, tags, context);
+        tagSqlDao.batchDeleteFromTransaction(objectId.toString(), objectType, existingTags, context);
+
+        tagAuditSqlDao.batchInsertFromTransaction(tags, context);
+        tagAuditSqlDao.batchDeleteFromTransaction(existingTags, context);
+    }
+
+    @Override
+    public void saveTagsFromTransaction(final Transmogrifier dao, final UUID objectId, final String objectType,
+                                        final List<Tag> tags, final CallContext context) {
         TagSqlDao tagSqlDao = dao.become(TagSqlDao.class);
 
         // get list of existing tags
@@ -72,6 +104,17 @@ public class AuditedTagDao implements TagDao {
     }
 
     @Override
+    public List<Tag> loadTags(final UUID objectId, final String objectType) {
+        return tagSqlDao.load(objectId.toString(), objectType);
+    }
+
+    @Override
+    public List<Tag> loadTagsFromTransaction(final Transmogrifier dao, final UUID objectId, final String objectType) {
+        TagSqlDao tagSqlDao = dao.become(TagSqlDao.class);
+        return tagSqlDao.load(objectId.toString(), objectType);
+    }
+
+    @Override
     public void addTag(final String tagName, final UUID objectId, final String objectType, final CallContext context) {
         tagSqlDao.inTransaction(new Transaction<Void, TagSqlDao>() {
             @Override

diff --git a/util/src/main/java/com/ning/billing/util/tag/dao/TagDao.java b/util/src/main/java/com/ning/billing/util/tag/dao/TagDao.java
index b4a8223..11594ac 100644
--- a/util/src/main/java/com/ning/billing/util/tag/dao/TagDao.java
+++ b/util/src/main/java/com/ning/billing/util/tag/dao/TagDao.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2010-2011 Ning, Inc.
+ *
+ * Ning 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 com.ning.billing.util.tag.dao;
 
 import com.ning.billing.util.callcontext.CallContext;
@@ -9,7 +25,13 @@ import java.util.List;
 import java.util.UUID;
 
 public interface TagDao {
-    void saveTags(Transmogrifier dao, UUID objectId, String objectType, List<Tag> tags, CallContext context);
+    void saveTagsFromTransaction(Transmogrifier dao, UUID objectId, String objectType, List<Tag> tags, CallContext context);
+
+    void saveTags(UUID objectId, String objectType, List<Tag> tags, CallContext context);
+
+    List<Tag> loadTags(UUID objectId, String objectType);
+
+    List<Tag> loadTagsFromTransaction(Transmogrifier dao, UUID objectId, String objectType);
 
     void addTag(String tagName, UUID objectId, String objectType, CallContext context);
 

diff --git a/util/src/main/java/com/ning/billing/util/tag/dao/TaggableDao.java b/util/src/main/java/com/ning/billing/util/tag/dao/TaggableDao.java
index f2e1f23..de45f87 100644
--- a/util/src/main/java/com/ning/billing/util/tag/dao/TaggableDao.java
+++ b/util/src/main/java/com/ning/billing/util/tag/dao/TaggableDao.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2010-2011 Ning, Inc.
+ *
+ * Ning 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 com.ning.billing.util.tag.dao;
 
 import com.ning.billing.util.callcontext.CallContext;

diff --git a/util/src/test/java/com/ning/billing/util/tag/dao/MockTagDao.java b/util/src/test/java/com/ning/billing/util/tag/dao/MockTagDao.java
new file mode 100644
index 0000000..ddce1cf
--- /dev/null
+++ b/util/src/test/java/com/ning/billing/util/tag/dao/MockTagDao.java
@@ -0,0 +1,104 @@
+/*
+ * Copyright 2010-2011 Ning, Inc.
+ *
+ * Ning 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 com.ning.billing.util.tag.dao;
+
+import com.ning.billing.util.callcontext.CallContext;
+import com.ning.billing.util.clock.Clock;
+import com.ning.billing.util.tag.Tag;
+import org.joda.time.DateTime;
+import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
+
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+public class MockTagDao implements TagDao {
+    private Map<UUID, List<Tag>> tagStore = new HashMap<UUID, List<Tag>>();
+    private final Clock clock;
+
+    public MockTagDao(Clock clock) {
+        this.clock = clock;
+    }
+
+    @Override
+    public void saveTagsFromTransaction(final Transmogrifier dao, final UUID objectId, final String objectType,
+                                        final List<Tag> tags, final CallContext context) {
+        tagStore.put(objectId, tags);
+    }
+
+    @Override
+    public void saveTags(UUID objectId, String objectType, List<Tag> tags, CallContext context) {
+        tagStore.put(objectId, tags);
+    }
+
+    @Override
+    public List<Tag> loadTags(UUID objectId, String objectType) {
+        return tagStore.get(objectId);
+    }
+
+    @Override
+    public List<Tag> loadTagsFromTransaction(Transmogrifier dao, UUID objectId, String objectType) {
+        return tagStore.get(objectId);
+    }
+
+    @Override
+    public void addTag(final String tagName, final UUID objectId, final String objectType, final CallContext context) {
+        Tag tag = new Tag() {
+            private UUID id = UUID.randomUUID();
+            private DateTime createdDate = clock.getUTCNow();
+
+            @Override
+            public String getTagDefinitionName() {
+                return tagName;
+            }
+
+            @Override
+            public UUID getId() {
+                return id;
+            }
+
+            @Override
+            public String getCreatedBy() {
+                return context.getUserName();
+            }
+
+            @Override
+            public DateTime getCreatedDate() {
+                return createdDate;
+            }
+        };
+
+
+        tagStore.get(objectId).add(tag);
+    }
+
+    @Override
+    public void removeTag(String tagName, UUID objectId, String objectType, CallContext context) {
+        List<Tag> tags = tagStore.get(objectId);
+        if (tags != null) {
+            Iterator<Tag> tagIterator = tags.iterator();
+            while (tagIterator.hasNext()) {
+                Tag tag = tagIterator.next();
+                if (tag.getTagDefinitionName().equals(tagName)) {
+                    tagIterator.remove();
+                }
+            }
+        }
+    }
+}

diff --git a/util/src/test/java/com/ning/billing/util/tag/MockTagStoreModuleMemory.java b/util/src/test/java/com/ning/billing/util/tag/MockTagStoreModuleMemory.java
new file mode 100644
index 0000000..3acc3a1
--- /dev/null
+++ b/util/src/test/java/com/ning/billing/util/tag/MockTagStoreModuleMemory.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2011 Ning, Inc.
+ *
+ * Ning 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 com.ning.billing.util.tag;
+
+import com.ning.billing.util.glue.TagStoreModule;
+import com.ning.billing.util.tag.dao.MockTagDao;
+import com.ning.billing.util.tag.dao.MockTagDefinitionDao;
+import com.ning.billing.util.tag.dao.TagDao;
+import com.ning.billing.util.tag.dao.TagDefinitionDao;
+
+public class MockTagStoreModuleMemory extends TagStoreModule {
+    @Override
+    protected void installDaos() {
+        bind(TagDefinitionDao.class).to(MockTagDefinitionDao.class).asEagerSingleton();
+        bind(TagDao.class).to(MockTagDao.class).asEagerSingleton();
+    }
+}

diff --git a/util/src/test/java/com/ning/billing/util/tag/TestTagStore.java b/util/src/test/java/com/ning/billing/util/tag/TestTagStore.java
index f742709..6213a2d 100644
--- a/util/src/test/java/com/ning/billing/util/tag/TestTagStore.java
+++ b/util/src/test/java/com/ning/billing/util/tag/TestTagStore.java
@@ -17,23 +17,23 @@
 package com.ning.billing.util.tag;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 
+import com.ning.billing.account.api.Account;
 import com.ning.billing.util.callcontext.CallContext;
 import com.ning.billing.util.callcontext.CallOrigin;
 import com.ning.billing.util.callcontext.UserType;
 import com.ning.billing.util.callcontext.DefaultCallContextFactory;
 import com.ning.billing.util.clock.Clock;
-import com.ning.billing.util.tag.dao.AuditedTagDao;
 import com.ning.billing.util.tag.dao.TagDao;
 import org.apache.commons.io.IOUtils;
 import org.joda.time.DateTime;
 import org.joda.time.Seconds;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.sqlobject.mixins.Transmogrifier;
 import org.skife.jdbi.v2.tweak.HandleCallback;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -47,7 +47,6 @@ import com.ning.billing.dbi.MysqlTestingHelper;
 
 import com.ning.billing.util.api.TagDefinitionApiException;
 import com.ning.billing.util.tag.dao.TagDefinitionDao;
-import com.ning.billing.util.tag.dao.TagSqlDao;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
@@ -57,27 +56,25 @@ import static org.testng.Assert.fail;
 
 
 @Test(groups={"slow"})
-@Guice(modules = TagStoreModuleMock.class)
+@Guice(modules = MockTagStoreModuleSql.class)
 public class TestTagStore {
-    private final static String ACCOUNT_TYPE = "ACCOUNT";
-
     @Inject
     private MysqlTestingHelper helper;
 
     @Inject
     private IDBI dbi;
 
-    private TagSqlDao tagSqlDao;
-
+    @Inject
+    private TagDao tagDao;
+    
     @Inject
     private TagDefinitionDao tagDefinitionDao;
 
     @Inject
     private Clock clock;
 
-    private TagDefinition tag1;
-    private TagDefinition tag2;
-
+    private TagDefinition testTag;
+    private TestSqlDao dao;
 
     private final Logger log = LoggerFactory.getLogger(TestTagStore.class);
     private CallContext context;
@@ -90,14 +87,13 @@ public class TestTagStore {
 
             helper.startMysql();
             helper.initDb(utilDdl);
-            tagSqlDao = dbi.onDemand(TagSqlDao.class);
-            tagSqlDao.test();
 
             context = new DefaultCallContextFactory(clock).createCallContext("Tag store test", CallOrigin.TEST, UserType.TEST);
-
+            dao = dbi.onDemand(TestSqlDao.class);
+            
             cleanupTags();
-            tag1 = tagDefinitionDao.create("tag1", "First tag", context);
-            tag2 = tagDefinitionDao.create("tag2", "Second tag", context);
+            tagDefinitionDao.create("tag1", "First tag", context);
+            testTag = tagDefinitionDao.create("testTag", "Second tag", context);
         }
         catch (Throwable t) {
             log.error("Failed to start tag store tests", t);
@@ -111,13 +107,6 @@ public class TestTagStore {
         helper.stopMysql();
     }
 
-    private void saveTags(final TagSqlDao dao, final String objectType, final UUID accountId,
-                          final List<Tag> tagList, final CallContext context)  {
-        TagDao tagDao = new AuditedTagDao(dbi);
-        tagDao.saveTags(dao, accountId, objectType, tagList, context);
-    }
-
-
     private void cleanupTags() {
         try {
             helper.getDBI().withHandle(new HandleCallback<Void>() {
@@ -137,14 +126,13 @@ public class TestTagStore {
     public void testTagCreationAndRetrieval() {
         UUID accountId = UUID.randomUUID();
 
-        TagStore tagStore = new DefaultTagStore(accountId, ACCOUNT_TYPE);
-        Tag tag = new DescriptiveTag(tag2);
+        TagStore tagStore = new DefaultTagStore(accountId, Account.ObjectType);
+        Tag tag = new DescriptiveTag(testTag);
         tagStore.add(tag);
 
-        TagSqlDao dao = dbi.onDemand(TagSqlDao.class);
-        saveTags(dao, ACCOUNT_TYPE, accountId, tagStore.getEntityList(), context);
+        tagDao.saveTagsFromTransaction(dao, accountId, Account.ObjectType, tagStore.getEntityList(), context);
 
-        List<Tag> savedTags = dao.load(accountId.toString(), ACCOUNT_TYPE);
+        List<Tag> savedTags = tagDao.loadTags(accountId, Account.ObjectType);
         assertEquals(savedTags.size(), 1);
 
         Tag savedTag = savedTags.get(0);
@@ -155,19 +143,19 @@ public class TestTagStore {
     @Test
     public void testControlTagCreation() {
         UUID accountId = UUID.randomUUID();
-        TagStore tagStore = new DefaultTagStore(accountId, ACCOUNT_TYPE);
+        TagStore tagStore = new DefaultTagStore(accountId, Account.ObjectType);
 
         ControlTag tag = new DefaultControlTag(ControlTagType.AUTO_INVOICING_OFF);
         tagStore.add(tag);
         assertEquals(tagStore.generateInvoice(), false);
 
         List<Tag> tagList = tagStore.getEntityList();
-        saveTags(tagSqlDao, ACCOUNT_TYPE, accountId, tagList, context);
+        tagDao.saveTagsFromTransaction(dao, accountId, Account.ObjectType, tagList, context);
 
         tagStore.clear();
         assertEquals(tagStore.getEntityList().size(), 0);
 
-        tagList = tagSqlDao.load(accountId.toString(), ACCOUNT_TYPE);
+        tagList = tagDao.loadTags(accountId, Account.ObjectType);
         tagStore.add(tagList);
         assertEquals(tagList.size(), 1);
 
@@ -177,7 +165,7 @@ public class TestTagStore {
     @Test
     public void testDescriptiveTagCreation() {
         UUID accountId = UUID.randomUUID();
-        TagStore tagStore = new DefaultTagStore(accountId, ACCOUNT_TYPE);
+        TagStore tagStore = new DefaultTagStore(accountId, Account.ObjectType);
 
         String definitionName = "SomeTestTag";
         TagDefinition tagDefinition = null;
@@ -191,13 +179,12 @@ public class TestTagStore {
         tagStore.add(tag);
         assertEquals(tagStore.generateInvoice(), true);
 
-        List<Tag> tagList = tagStore.getEntityList();
-        saveTags(tagSqlDao, ACCOUNT_TYPE, accountId, tagList, context);
+        tagDao.saveTagsFromTransaction(dao, accountId, Account.ObjectType, tagStore.getEntityList(), context);
 
         tagStore.clear();
         assertEquals(tagStore.getEntityList().size(), 0);
 
-        tagList = tagSqlDao.load(accountId.toString(), ACCOUNT_TYPE);
+        List<Tag> tagList = tagDao.loadTags(accountId, Account.ObjectType);
         tagStore.add(tagList);
         assertEquals(tagList.size(), 1);
 
@@ -207,7 +194,7 @@ public class TestTagStore {
     @Test
     public void testMixedTagCreation() {
         UUID accountId = UUID.randomUUID();
-        TagStore tagStore = new DefaultTagStore(accountId, ACCOUNT_TYPE);
+        TagStore tagStore = new DefaultTagStore(accountId, Account.ObjectType);
 
         String definitionName = "MixedTagTest";
         TagDefinition tagDefinition = null;
@@ -225,13 +212,12 @@ public class TestTagStore {
         tagStore.add(controlTag);
         assertEquals(tagStore.generateInvoice(), false);
 
-        List<Tag> tagList = tagStore.getEntityList();
-        saveTags(tagSqlDao, ACCOUNT_TYPE, accountId, tagList, context);
+        tagDao.saveTagsFromTransaction(dao, accountId, Account.ObjectType, tagStore.getEntityList(), context);
 
         tagStore.clear();
         assertEquals(tagStore.getEntityList().size(), 0);
 
-        tagList = tagSqlDao.load(accountId.toString(), ACCOUNT_TYPE);
+        List<Tag> tagList = tagDao.loadTags(accountId, Account.ObjectType);
         tagStore.add(tagList);
         assertEquals(tagList.size(), 2);
 
@@ -241,7 +227,7 @@ public class TestTagStore {
     @Test
     public void testControlTags() {
         UUID accountId = UUID.randomUUID();
-        TagStore tagStore = new DefaultTagStore(accountId, ACCOUNT_TYPE);
+        TagStore tagStore = new DefaultTagStore(accountId, Account.ObjectType);
         assertEquals(tagStore.generateInvoice(), true);
         assertEquals(tagStore.processPayment(), true);
 
@@ -284,14 +270,13 @@ public class TestTagStore {
         assertNotNull(tagDefinition);
 
         UUID objectId = UUID.randomUUID();
-        String objectType = "TestType";
-        TagStore tagStore = new DefaultTagStore(objectId, objectType);
+        TagStore tagStore = new DefaultTagStore(objectId, Account.ObjectType);
         Tag tag = new DescriptiveTag(tagDefinition);
         tagStore.add(tag);
 
-        saveTags(tagSqlDao, objectType, objectId, tagStore.getEntityList(), context);
+        tagDao.saveTags(objectId, Account.ObjectType, tagStore.getEntityList(), context);
 
-        List<Tag> tags = tagSqlDao.load(objectId.toString(), objectType);
+        List<Tag> tags = tagDao.loadTags(objectId, Account.ObjectType);
         assertEquals(tags.size(), 1);
 
         tagDefinitionDao.deleteTagDefinition(definitionName, context);
@@ -310,14 +295,13 @@ public class TestTagStore {
         assertNotNull(tagDefinition);
 
         UUID objectId = UUID.randomUUID();
-        String objectType = "TestType";
-        TagStore tagStore = new DefaultTagStore(objectId, objectType);
+        TagStore tagStore = new DefaultTagStore(objectId, Account.ObjectType);
         Tag tag = new DescriptiveTag(tagDefinition);
         tagStore.add(tag);
 
-        saveTags(tagSqlDao, objectType, objectId, tagStore.getEntityList(), context);
+        tagDao.saveTags(objectId, Account.ObjectType, tagStore.getEntityList(), context);
 
-        List<Tag> tags = tagSqlDao.load(objectId.toString(), objectType);
+        List<Tag> tags = tagDao.loadTags(objectId, Account.ObjectType);
         assertEquals(tags.size(), 1);
 
         try {
@@ -390,14 +374,13 @@ public class TestTagStore {
     public void testTagInsertAudit() {
         UUID accountId = UUID.randomUUID();
 
-        TagStore tagStore = new DefaultTagStore(accountId, ACCOUNT_TYPE);
-        Tag tag = new DescriptiveTag(tag2);
+        TagStore tagStore = new DefaultTagStore(accountId, Account.ObjectType);
+        Tag tag = new DescriptiveTag(testTag);
         tagStore.add(tag);
 
-        TagSqlDao dao = dbi.onDemand(TagSqlDao.class);
-        saveTags(dao, ACCOUNT_TYPE, accountId, tagStore.getEntityList(), context);
+        tagDao.saveTagsFromTransaction(dao, accountId, Account.ObjectType, tagStore.getEntityList(), context);
 
-        List<Tag> savedTags = dao.load(accountId.toString(), ACCOUNT_TYPE);
+        List<Tag> savedTags = tagDao.loadTags(accountId, Account.ObjectType);
         assertEquals(savedTags.size(), 1);
 
         Tag savedTag = savedTags.get(0);
@@ -421,19 +404,20 @@ public class TestTagStore {
     public void testTagDeleteAudit() {
         UUID accountId = UUID.randomUUID();
 
-        TagStore tagStore = new DefaultTagStore(accountId, ACCOUNT_TYPE);
-        Tag tag = new DescriptiveTag(tag2);
+        TagStore tagStore = new DefaultTagStore(accountId, Account.ObjectType);
+        Tag tag = new DescriptiveTag(testTag);
         tagStore.add(tag);
 
-        TagSqlDao dao = dbi.onDemand(TagSqlDao.class);
-        saveTags(dao, ACCOUNT_TYPE, accountId, tagStore.getEntityList(), context);
-        saveTags(dao, ACCOUNT_TYPE, accountId, new ArrayList<Tag>(), context);
+        tagDao.saveTags(accountId, Account.ObjectType, tagStore.getEntityList(), context);
+
+        tagStore.remove(tag);
+        tagDao.saveTags(accountId, Account.ObjectType, tagStore.getEntityList(), context);
 
-        List<Tag> savedTags = dao.load(accountId.toString(), ACCOUNT_TYPE);
+        List<Tag> savedTags = tagDao.loadTags(accountId, Account.ObjectType);
         assertEquals(savedTags.size(), 0);
 
         Handle handle = dbi.open();
-        String query = String.format("select * from audit_log where table_name = 'Tag' and record_id='%s' and change_type='DELETE'",
+        String query = String.format("select * from audit_log where table_name = 'tag' and record_id='%s' and change_type='DELETE'",
                                      tag.getId().toString());
         List<Map<String, Object>> result = handle.select(query);
         assertNotNull(result);
@@ -443,22 +427,6 @@ public class TestTagStore {
         assertTrue(Seconds.secondsBetween(changeDate, context.getUpdatedDate()).getSeconds() < 2);
         assertEquals(result.get(0).get("changed_by"), context.getUserName());
     }
-
-//    @Test
-//    public void testTagDefinitionInsertAudit() {
-//
-//        fail();
-//    }
-//
-//    @Test
-//    public void testTagDefinitionUpdateAudit() {
-//
-//        fail();
-//    }
-//
-//    @Test
-//    public void testTagDefinitionDeleteAudit() {
-//
-//        fail();
-//    }
+    
+    public interface TestSqlDao extends Transmogrifier {}
 }