@Test
  public void testUpdateDataUsage() throws PersistenceException {
    ArgumentCaptor<String> keyArg1 = ArgumentCaptor.forClass(String.class);
    ArgumentCaptor<String> keyArg2 = ArgumentCaptor.forClass(String.class);
    ArgumentCaptor<String> cqlArg = ArgumentCaptor.forClass(String.class);
    ArgumentCaptor<PersistentItem> itemArg = ArgumentCaptor.forClass(PersistentItem.class);

    attributesList = new ArrayList<>();
    Map<String, Object> attributes = new HashMap<>();
    attributes.put(DATA_USAGE_LONG, LONG_1);
    attributes.put(DATA_LIMIT_LONG, LONG_1);
    attributesList.add(attributes);
    when(persistentStore.get(anyString(), anyString())).thenReturn(attributesList);

    attributesStore.updateUserDataUsage(USER, LONG_5);

    verify(persistentStore, atLeast(2)).get(keyArg1.capture(), cqlArg.capture());
    verify(persistentStore).add(keyArg2.capture(), itemArg.capture());

    assertThat(keyArg1.getValue(), is(PersistentStore.USER_ATTRIBUTE_TYPE));
    assertThat(keyArg2.getValue(), is(PersistentStore.USER_ATTRIBUTE_TYPE));

    assertThat(itemArg.getValue().getLongProperty(AttributesStore.DATA_USAGE_KEY), is(600L));

    assertThat(cqlArg.getValue(), is(CQL));
  }
 @Test
 public void testGetDataLimit() throws PersistenceException {
   attributesList = new ArrayList<>();
   Map<String, Object> attributes = new HashMap<>();
   attributes.put(DATA_USAGE_LONG, LONG_2);
   attributes.put(DATA_LIMIT_LONG, LONG_1);
   attributesList.add(attributes);
   when(persistentStore.get(anyString(), anyString())).thenReturn(attributesList);
   long usage = attributesStore.getDataLimitByUser(USER);
   assertThat(usage, is(LONG_1));
 }
 private List<Map<String, Object>> query(String q) throws PersistenceException {
   LOCK.lock();
   try {
     List<Map<String, Object>> results =
         persistentStore.get(SubscriptionsPersistentStore.SUBSCRIPTIONS_TYPE, q);
     assert results.size() <= 1;
     return results;
   } finally {
     LOCK.unlock();
   }
 }
  @Test
  public void resetUserDataUsages() throws PersistenceException {
    attributesList = new ArrayList<>();
    Map<String, Object> attributes = new HashMap<>();
    attributes.put(DATA_USAGE_LONG, LONG_2);
    attributes.put(DATA_LIMIT_LONG, LONG_1);
    attributesList.add(attributes);
    when(persistentStore.get(anyString())).thenReturn(attributesList);

    ArgumentCaptor<String> keyArg = ArgumentCaptor.forClass(String.class);
    ArgumentCaptor<PersistentItem> itemArg = ArgumentCaptor.forClass(PersistentItem.class);
    attributesStore.resetUserDataUsages();
    verify(persistentStore).add(keyArg.capture(), itemArg.capture());
    assertThat(keyArg.getValue(), is(PersistentStore.USER_ATTRIBUTE_TYPE));

    assertThat(itemArg.getValue().getLongProperty(AttributesStore.DATA_USAGE_KEY), is(0L));
    assertThat(
        itemArg.getValue().getLongProperty(AttributesStore.DATA_USAGE_LIMIT_KEY), is(LONG_1));
  }
 @Test(expected = PersistenceException.class)
 public void testGetAllUsersThrowsException() throws PersistenceException {
   when(persistentStore.get(anyString())).thenThrow(new PersistenceException());
   attributesStore.getAllUsers();
 }
  @Test(expected = PersistenceException.class)
  public void testPersistenceStoreThrowsExceptionOnGet() throws PersistenceException {
    when(persistentStore.get(anyString(), anyString())).thenThrow(new PersistenceException());

    attributesStore.updateUserDataUsage(USER, LONG_5);
  }
 @Test
 public void testPersistenceStoreReturnsEmptyList() throws PersistenceException {
   when(persistentStore.get(anyString(), anyString())).thenReturn(new ArrayList<>());
   assertThat(attributesStore.getCurrentDataUsageByUser(USER), is(0L));
 }