killbill-aplcache

diff --git a/account/src/test/java/com/ning/billing/account/api/user/TestDefaultAccountUserApi.java b/account/src/test/java/com/ning/billing/account/api/user/TestDefaultAccountUserApi.java
index 51111aa..9dd5f33 100644
--- a/account/src/test/java/com/ning/billing/account/api/user/TestDefaultAccountUserApi.java
+++ b/account/src/test/java/com/ning/billing/account/api/user/TestDefaultAccountUserApi.java
@@ -27,6 +27,7 @@ import org.testng.annotations.Test;
 import com.ning.billing.account.AccountTestSuite;
 import com.ning.billing.account.api.Account;
 import com.ning.billing.account.api.AccountData;
+import com.ning.billing.account.api.AccountEmail;
 import com.ning.billing.account.api.BillCycleDay;
 import com.ning.billing.account.api.DefaultAccount;
 import com.ning.billing.account.api.DefaultAccountEmail;
@@ -36,11 +37,11 @@ import com.ning.billing.account.dao.AccountEmailDao;
 import com.ning.billing.account.dao.MockAccountDao;
 import com.ning.billing.account.dao.MockAccountEmailDao;
 import com.ning.billing.catalog.api.Currency;
-import com.ning.billing.util.svcsapi.bus.InternalBus;
 import com.ning.billing.util.callcontext.CallContext;
 import com.ning.billing.util.callcontext.CallContextFactory;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.callcontext.InternalTenantContext;
+import com.ning.billing.util.svcsapi.bus.InternalBus;
 
 public class TestDefaultAccountUserApi extends AccountTestSuite {
 
@@ -118,17 +119,19 @@ public class TestDefaultAccountUserApi extends AccountTestSuite {
         Assert.assertEquals(accountEmailDao.getByAccountId(accountId, tenantContext).size(), 0);
 
         // Add the first email
-        final String email1 = UUID.randomUUID().toString();
-        accountUserApi.addEmail(accountId, new DefaultAccountEmail(accountId, email1), callContext);
+        final String emailAddress1 = UUID.randomUUID().toString();
+        final AccountEmail email1 = new DefaultAccountEmail(accountId, emailAddress1);
+        accountUserApi.addEmail(accountId, email1, callContext);
         Assert.assertEquals(accountEmailDao.getByAccountId(accountId, tenantContext).size(), 1);
 
         // Add a second one
-        final String email2 = UUID.randomUUID().toString();
-        accountUserApi.addEmail(accountId, new DefaultAccountEmail(accountId, email2), callContext);
+        final String emailAddress2 = UUID.randomUUID().toString();
+        final AccountEmail email2 = new DefaultAccountEmail(accountId, emailAddress2);
+        accountUserApi.addEmail(accountId, email2, callContext);
         Assert.assertEquals(accountEmailDao.getByAccountId(accountId, tenantContext).size(), 2);
 
         // Remove the first second one
-        accountUserApi.removeEmail(accountId, new DefaultAccountEmail(accountId, email1), callContext);
+        accountUserApi.removeEmail(accountId, email1, callContext);
         Assert.assertEquals(accountEmailDao.getByAccountId(accountId, tenantContext).size(), 1);
     }
 }

diff --git a/account/src/test/java/com/ning/billing/account/dao/MockAccountEmailDao.java b/account/src/test/java/com/ning/billing/account/dao/MockAccountEmailDao.java
index aa4c401..f408592 100644
--- a/account/src/test/java/com/ning/billing/account/dao/MockAccountEmailDao.java
+++ b/account/src/test/java/com/ning/billing/account/dao/MockAccountEmailDao.java
@@ -16,69 +16,28 @@
 
 package com.ning.billing.account.dao;
 
-import java.util.HashSet;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
 
 import com.ning.billing.account.api.AccountApiException;
 import com.ning.billing.account.api.AccountEmail;
-import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalTenantContext;
-import com.ning.billing.util.entity.EntityPersistenceException;
+import com.ning.billing.util.entity.dao.MockEntityDaoBase;
 
-import com.google.common.collect.ImmutableList;
-
-public class MockAccountEmailDao implements AccountEmailDao {
-
-    private final Map<UUID, Set<AccountEmail>> emails = new ConcurrentHashMap<UUID, Set<AccountEmail>>();
-
-    @Override
-    public void create(final AccountEmail entity, final InternalCallContext context) throws AccountApiException {
-        Set<AccountEmail> theSet = emails.get(entity.getAccountId());
-        theSet.add(entity);
-    }
-
-    @Override
-    public Long getRecordId(final UUID id, final InternalTenantContext context) {
-        return 1L;
-    }
-
-    @Override
-    public AccountEmail getById(final UUID id, final InternalTenantContext context) {
-        throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public List<AccountEmail> get(final InternalTenantContext context) {
-        throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public void delete(final AccountEmail email, final InternalCallContext context) {
-        Set<AccountEmail> theSet = emails.get(email.getAccountId());
-        theSet.remove(email);
-    }
+public class MockAccountEmailDao extends MockEntityDaoBase<AccountEmail, AccountApiException> implements AccountEmailDao {
 
     @Override
     public List<AccountEmail> getByAccountId(final UUID accountId, final InternalTenantContext context) {
-        final Set<AccountEmail> accountEmails = emails.get(accountId);
-        if (accountEmails == null) {
-            return ImmutableList.<AccountEmail>of();
-        } else {
-            return ImmutableList.<AccountEmail>copyOf(accountEmails.iterator());
+        final List<AccountEmail> accountEmails = new ArrayList<AccountEmail>();
+        for (final Map<Long, AccountEmail> accountEmail : entities.values()) {
+            final AccountEmail email = accountEmail.values().iterator().next();
+            if (email.getAccountId().equals(accountId)) {
+                accountEmails.add(email);
+            }
         }
-    }
 
-    @Override
-    public void test(final InternalTenantContext context) {
-    }
-
-    @Override
-    public AccountEmail getByRecordId(final Long recordId,
-            final InternalTenantContext context) {
-        return null;
+        return accountEmails;
     }
 }