util/kvstore/InMemoryStoreSuite.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.util.kvstore;
0019
0020 import java.util.NoSuchElementException;
0021
0022 import com.google.common.collect.ImmutableSet;
0023 import org.junit.Test;
0024 import static org.junit.Assert.*;
0025
0026 public class InMemoryStoreSuite {
0027
0028   @Test
0029   public void testObjectWriteReadDelete() throws Exception {
0030     KVStore store = new InMemoryStore();
0031
0032     CustomType1 t = new CustomType1();
0033     t.key = "key";
0034     t.id = "id";
0035     t.name = "name";
0036
0037     try {
0038       store.read(CustomType1.class, t.key);
0039       fail("Expected exception for non-existent object.");
0040     } catch (NoSuchElementException nsee) {
0041       // Expected.
0042     }
0043
0044     store.write(t);
0045     assertEquals(t, store.read(t.getClass(), t.key));
0046     assertEquals(1L, store.count(t.getClass()));
0047
0048     store.delete(t.getClass(), t.key);
0049     try {
0050       store.read(t.getClass(), t.key);
0051       fail("Expected exception for deleted object.");
0052     } catch (NoSuchElementException nsee) {
0053       // Expected.
0054     }
0055   }
0056
0057   @Test
0058   public void testMultipleObjectWriteReadDelete() throws Exception {
0059     KVStore store = new InMemoryStore();
0060
0061     CustomType1 t1 = new CustomType1();
0062     t1.key = "key1";
0063     t1.id = "id";
0064     t1.name = "name1";
0065
0066     CustomType1 t2 = new CustomType1();
0067     t2.key = "key2";
0068     t2.id = "id";
0069     t2.name = "name2";
0070
0071     store.write(t1);
0072     store.write(t2);
0073
0074     assertEquals(t1, store.read(t1.getClass(), t1.key));
0075     assertEquals(t2, store.read(t2.getClass(), t2.key));
0076     assertEquals(2L, store.count(t1.getClass()));
0077
0078     store.delete(t1.getClass(), t1.key);
0079     assertEquals(t2, store.read(t2.getClass(), t2.key));
0080     store.delete(t2.getClass(), t2.key);
0081     try {
0082       store.read(t2.getClass(), t2.key);
0083       fail("Expected exception for deleted object.");
0084     } catch (NoSuchElementException nsee) {
0085       // Expected.
0086     }
0087   }
0088
0089   @Test
0090   public void testMetadata() throws Exception {
0091     KVStore store = new InMemoryStore();
0092     assertNull(store.getMetadata(CustomType1.class));
0093
0094     CustomType1 t = new CustomType1();
0095     t.id = "id";
0096     t.name = "name";
0097
0098     store.setMetadata(t);
0099     assertEquals(t, store.getMetadata(CustomType1.class));
0100
0101     store.setMetadata(null);
0102     assertNull(store.getMetadata(CustomType1.class));
0103   }
0104
0105   @Test
0106   public void testUpdate() throws Exception {
0107     KVStore store = new InMemoryStore();
0108
0109     CustomType1 t = new CustomType1();
0110     t.key = "key";
0111     t.id = "id";
0112     t.name = "name";
0113
0114     store.write(t);
0115
0116     t.name = "anotherName";
0117
0118     store.write(t);
0119     assertEquals(1, store.count(t.getClass()));
0120     assertSame(t, store.read(t.getClass(), t.key));
0121   }
0122
0123   @Test
0124   public void testArrayIndices() throws Exception {
0125     KVStore store = new InMemoryStore();
0126
0127     ArrayKeyIndexType o = new ArrayKeyIndexType();
0128     o.key = new int[] { 1, 2 };
0129     o.id = new String[] { "3", "4" };
0130
0131     store.write(o);
0132     assertEquals(o, store.read(ArrayKeyIndexType.class, o.key));
0133     assertEquals(o, store.view(ArrayKeyIndexType.class).index("id").first(o.id).iterator().next());
0134   }
0135
0136   @Test
0137   public void testRemoveAll() throws Exception {
0138     KVStore store = new InMemoryStore();
0139
0140     for (int i = 0; i < 2; i++) {
0141       for (int j = 0; j < 2; j++) {
0142         ArrayKeyIndexType o = new ArrayKeyIndexType();
0143         o.key = new int[] { i, j, 0 };
0144         o.id = new String[] { "things" };
0145         store.write(o);
0146
0147         o = new ArrayKeyIndexType();
0148         o.key = new int[] { i, j, 1 };
0149         o.id = new String[] { "more things" };
0150         store.write(o);
0151       }
0152     }
0153
0154     ArrayKeyIndexType o = new ArrayKeyIndexType();
0155     o.key = new int[] { 2, 2, 2 };
0156     o.id = new String[] { "things" };
0157     store.write(o);
0158
0159     assertEquals(9, store.count(ArrayKeyIndexType.class));
0160
0161     // Try removing non-existing keys
0162     assert(!store.removeAllByIndexValues(
0163       ArrayKeyIndexType.class,
0164       KVIndex.NATURAL_INDEX_NAME,
0165       ImmutableSet.of(new int[] {10, 10, 10}, new int[] { 3, 3, 3 })));
0166     assertEquals(9, store.count(ArrayKeyIndexType.class));
0167
0168     assert(store.removeAllByIndexValues(
0169       ArrayKeyIndexType.class,
0170       KVIndex.NATURAL_INDEX_NAME,
0171       ImmutableSet.of(new int[] {0, 0, 0}, new int[] { 2, 2, 2 })));
0172     assertEquals(7, store.count(ArrayKeyIndexType.class));
0173
0174     assert(store.removeAllByIndexValues(
0175       ArrayKeyIndexType.class,
0176       "id",
0177       ImmutableSet.of(new String [] { "things" })));
0178     assertEquals(4, store.count(ArrayKeyIndexType.class));
0179
0180     assert(store.removeAllByIndexValues(
0181       ArrayKeyIndexType.class,
0182       "id",
0183       ImmutableSet.of(new String [] { "more things" })));
0184     assertEquals(0, store.count(ArrayKeyIndexType.class));
0185   }
0186
0187   @Test
0188   public void testBasicIteration() throws Exception {
0189     KVStore store = new InMemoryStore();
0190
0191     CustomType1 t1 = new CustomType1();
0192     t1.key = "1";
0193     t1.id = "id1";
0194     t1.name = "name1";
0195     store.write(t1);
0196
0197     CustomType1 t2 = new CustomType1();
0198     t2.key = "2";
0199     t2.id = "id2";
0200     t2.name = "name2";
0201     store.write(t2);
0202
0203     assertEquals(t1.id, store.view(t1.getClass()).iterator().next().id);
0204     assertEquals(t2.id, store.view(t1.getClass()).skip(1).iterator().next().id);
0205     assertEquals(t2.id, store.view(t1.getClass()).skip(1).max(1).iterator().next().id);
0206     assertEquals(t1.id,
0207       store.view(t1.getClass()).first(t1.key).max(1).iterator().next().id);
0208     assertEquals(t2.id,
0209       store.view(t1.getClass()).first(t2.key).max(1).iterator().next().id);
0210     assertFalse(store.view(t1.getClass()).first(t2.id).skip(1).iterator().hasNext());
0211   }
0212
0213   @Test
0214   public void testDeleteParentIndex() throws Exception {
0215     KVStore store = new InMemoryStore();
0216
0217     CustomType2 t1 = new CustomType2();
0218     t1.key = "key1";
0219     t1.id = "id1";
0220     t1.parentId = "parentId1";
0221     store.write(t1);
0222
0223     CustomType2 t2 = new CustomType2();
0224     t2.key = "key2";
0225     t2.id = "id2";
0226     t2.parentId = "parentId1";
0227     store.write(t2);
0228
0229     CustomType2 t3 = new CustomType2();
0230     t3.key = "key3";
0231     t3.id = "id1";
0232     t3.parentId = "parentId2";
0233     store.write(t3);
0234
0235     CustomType2 t4 = new CustomType2();
0236     t4.key = "key4";
0237     t4.id = "id2";
0238     t4.parentId = "parentId2";
0239     store.write(t4);
0240
0241     assertEquals(4, store.count(CustomType2.class));
0242
0243     store.delete(t1.getClass(), t1.key);
0244     assertEquals(3, store.count(CustomType2.class));
0245
0246     store.delete(t2.getClass(), t2.key);
0247     assertEquals(2, store.count(CustomType2.class));
0248
0249     store.delete(t3.getClass(), t3.key);
0250     assertEquals(1, store.count(CustomType2.class));
0251
0252     store.delete(t4.getClass(), t4.key);
0253     assertEquals(0, store.count(CustomType2.class));
0254   }
0255 }