Move in memory data structures into project.

10 files changed, 105 insertions, 52 deletions

diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/InMemoryAxolotlStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryAxolotlStore.java
index f4bf38eb..4e3c818d 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/InMemoryAxolotlStore.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryAxolotlStore.java
@@ -1,5 +1,9 @@
-package org.whispersystems.libaxolotl;
+package org.whispersystems.libaxolotl.state.impl;
 
+import org.whispersystems.libaxolotl.AxolotlAddress;
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.libaxolotl.state.AxolotlStore;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.state.SessionRecord;
@@ -9,11 +13,15 @@ import java.util.List;
 
 public class InMemoryAxolotlStore implements AxolotlStore {
 
-  private final InMemoryIdentityKeyStore  identityKeyStore  = new InMemoryIdentityKeyStore();
   private final InMemoryPreKeyStore       preKeyStore       = new InMemoryPreKeyStore();
   private final InMemorySessionStore      sessionStore      = new InMemorySessionStore();
   private final InMemorySignedPreKeyStore signedPreKeyStore = new InMemorySignedPreKeyStore();
 
+  private final InMemoryIdentityKeyStore  identityKeyStore;
+
+  public InMemoryAxolotlStore(IdentityKeyPair identityKeyPair, int registrationId) {
+    this.identityKeyStore = new InMemoryIdentityKeyStore(identityKeyPair, registrationId);
+  }
 
   @Override
   public IdentityKeyPair getIdentityKeyPair() {
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/InMemoryIdentityKeyStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryIdentityKeyStore.java
index bf6dc318..b1b04e08 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/InMemoryIdentityKeyStore.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryIdentityKeyStore.java
@@ -1,5 +1,7 @@
-package org.whispersystems.libaxolotl;
+package org.whispersystems.libaxolotl.state.impl;
 
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
 import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
 import org.whispersystems.libaxolotl.state.IdentityKeyStore;
@@ -16,16 +18,9 @@ public class InMemoryIdentityKeyStore implements IdentityKeyStore {
   private final IdentityKeyPair identityKeyPair;
   private final int             localRegistrationId;
 
-  public InMemoryIdentityKeyStore() {
-    try {
-      ECKeyPair identityKeyPairKeys = Curve.generateKeyPair();
-
-      this.identityKeyPair = new IdentityKeyPair(new IdentityKey(identityKeyPairKeys.getPublicKey()),
-                                                 identityKeyPairKeys.getPrivateKey());
-      this.localRegistrationId = SecureRandom.getInstance("SHA1PRNG").nextInt(16380) + 1;
-    } catch (NoSuchAlgorithmException e) {
-      throw new AssertionError(e);
-    }
+  public InMemoryIdentityKeyStore(IdentityKeyPair identityKeyPair, int localRegistrationId) {
+    this.identityKeyPair     = identityKeyPair;
+    this.localRegistrationId = localRegistrationId;
   }
 
   @Override
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/InMemoryPreKeyStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryPreKeyStore.java
index a2ea6811..80651f21 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/InMemoryPreKeyStore.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryPreKeyStore.java
@@ -1,5 +1,6 @@
-package org.whispersystems.libaxolotl;
+package org.whispersystems.libaxolotl.state.impl;
 
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/InMemorySessionStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySessionStore.java
index f707773e..f5270c7d 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/InMemorySessionStore.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySessionStore.java
@@ -1,5 +1,6 @@
-package org.whispersystems.libaxolotl;
+package org.whispersystems.libaxolotl.state.impl;
 
+import org.whispersystems.libaxolotl.AxolotlAddress;
 import org.whispersystems.libaxolotl.state.SessionRecord;
 import org.whispersystems.libaxolotl.state.SessionStore;
 
@@ -33,7 +34,9 @@ public class InMemorySessionStore implements SessionStore {
     List<Integer> deviceIds = new LinkedList<>();
 
     for (AxolotlAddress key : sessions.keySet()) {
-      if (key.getName().equals(name)) {
+      if (key.getName().equals(name) &&
+          key.getDeviceId() != 1)
+      {
         deviceIds.add(key.getDeviceId());
       }
     }
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/InMemorySignedPreKeyStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySignedPreKeyStore.java
index 9f452d6d..bab83137 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/InMemorySignedPreKeyStore.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySignedPreKeyStore.java
@@ -1,5 +1,6 @@
-package org.whispersystems.libaxolotl;
+package org.whispersystems.libaxolotl.state.impl;
 
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.libaxolotl.state.SignedPreKeyRecord;
 import org.whispersystems.libaxolotl.state.SignedPreKeyStore;
 
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java b/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java
index 539131f0..1105ffc6 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java
+++ b/tests/src/test/java/org/whispersystems/libaxolotl/SessionBuilderTest.java
@@ -25,10 +25,10 @@ public class SessionBuilderTest extends TestCase {
 
   public void testBasicPreKeyV2()
       throws InvalidKeyException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException, UntrustedIdentityException, NoSessionException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    AxolotlStore bobStore      = new InMemoryAxolotlStore();
+    AxolotlStore bobStore      = new TestInMemoryAxolotlStore();
     ECKeyPair    bobPreKeyPair = Curve.generateKeyPair();
     PreKeyBundle bobPreKey     = new PreKeyBundle(bobStore.getLocalRegistrationId(), 1,
                                                   31337, bobPreKeyPair.getPublicKey(),
@@ -64,7 +64,7 @@ public class SessionBuilderTest extends TestCase {
 
     runInteraction(aliceStore, bobStore);
 
-    aliceStore          = new InMemoryAxolotlStore();
+    aliceStore          = new TestInMemoryAxolotlStore();
     aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
     aliceSessionCipher  = new SessionCipher(aliceStore, BOB_ADDRESS);
 
@@ -104,10 +104,10 @@ public class SessionBuilderTest extends TestCase {
 
   public void testBasicPreKeyV3()
       throws InvalidKeyException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException, UntrustedIdentityException, NoSessionException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    final AxolotlStore bobStore                 = new InMemoryAxolotlStore();
+    final AxolotlStore bobStore                 = new TestInMemoryAxolotlStore();
           ECKeyPair    bobPreKeyPair            = Curve.generateKeyPair();
           ECKeyPair    bobSignedPreKeyPair      = Curve.generateKeyPair();
           byte[]       bobSignedPreKeySignature = Curve.calculateSignature(bobStore.getIdentityKeyPair().getPrivateKey(),
@@ -156,7 +156,7 @@ public class SessionBuilderTest extends TestCase {
 
     runInteraction(aliceStore, bobStore);
 
-    aliceStore          = new InMemoryAxolotlStore();
+    aliceStore          = new TestInMemoryAxolotlStore();
     aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
     aliceSessionCipher  = new SessionCipher(aliceStore, BOB_ADDRESS);
 
@@ -198,10 +198,10 @@ public class SessionBuilderTest extends TestCase {
   }
 
   public void testBadSignedPreKeySignature() throws InvalidKeyException, UntrustedIdentityException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
+    IdentityKeyStore bobIdentityKeyStore = new TestInMemoryIdentityKeyStore();
 
     ECKeyPair bobPreKeyPair            = Curve.generateKeyPair();
     ECKeyPair bobSignedPreKeyPair      = Curve.generateKeyPair();
@@ -237,10 +237,10 @@ public class SessionBuilderTest extends TestCase {
   }
 
   public void testRepeatBundleMessageV2() throws InvalidKeyException, UntrustedIdentityException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException, NoSessionException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    AxolotlStore bobStore = new InMemoryAxolotlStore();
+    AxolotlStore bobStore = new TestInMemoryAxolotlStore();
 
     ECKeyPair bobPreKeyPair            = Curve.generateKeyPair();
     ECKeyPair bobSignedPreKeyPair      = Curve.generateKeyPair();
@@ -290,10 +290,10 @@ public class SessionBuilderTest extends TestCase {
   }
 
   public void testRepeatBundleMessageV3() throws InvalidKeyException, UntrustedIdentityException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException, NoSessionException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    AxolotlStore bobStore = new InMemoryAxolotlStore();
+    AxolotlStore bobStore = new TestInMemoryAxolotlStore();
 
     ECKeyPair bobPreKeyPair            = Curve.generateKeyPair();
     ECKeyPair bobSignedPreKeyPair      = Curve.generateKeyPair();
@@ -344,10 +344,10 @@ public class SessionBuilderTest extends TestCase {
   }
 
   public void testBadMessageBundle() throws InvalidKeyException, UntrustedIdentityException, InvalidVersionException, InvalidMessageException, DuplicateMessageException, LegacyMessageException, InvalidKeyIdException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    AxolotlStore bobStore = new InMemoryAxolotlStore();
+    AxolotlStore bobStore = new TestInMemoryAxolotlStore();
 
     ECKeyPair bobPreKeyPair            = Curve.generateKeyPair();
     ECKeyPair bobSignedPreKeyPair      = Curve.generateKeyPair();
@@ -397,10 +397,10 @@ public class SessionBuilderTest extends TestCase {
   }
 
   public void testBasicKeyExchange() throws InvalidKeyException, LegacyMessageException, InvalidMessageException, DuplicateMessageException, UntrustedIdentityException, StaleKeyExchangeException, InvalidVersionException, NoSessionException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    AxolotlStore   bobStore          = new InMemoryAxolotlStore();
+    AxolotlStore   bobStore          = new TestInMemoryAxolotlStore();
     SessionBuilder bobSessionBuilder = new SessionBuilder(bobStore, ALICE_ADDRESS);
 
     KeyExchangeMessage aliceKeyExchangeMessage      = aliceSessionBuilder.process();
@@ -420,7 +420,7 @@ public class SessionBuilderTest extends TestCase {
 
     runInteraction(aliceStore, bobStore);
 
-    aliceStore              = new InMemoryAxolotlStore();
+    aliceStore              = new TestInMemoryAxolotlStore();
     aliceSessionBuilder     = new SessionBuilder(aliceStore, BOB_ADDRESS);
     aliceKeyExchangeMessage = aliceSessionBuilder.process();
 
@@ -439,10 +439,10 @@ public class SessionBuilderTest extends TestCase {
 
   public void testSimultaneousKeyExchange()
       throws InvalidKeyException, DuplicateMessageException, LegacyMessageException, InvalidMessageException, UntrustedIdentityException, StaleKeyExchangeException, NoSessionException {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    AxolotlStore   bobStore          = new InMemoryAxolotlStore();
+    AxolotlStore   bobStore          = new TestInMemoryAxolotlStore();
     SessionBuilder bobSessionBuilder = new SessionBuilder(bobStore, ALICE_ADDRESS);
 
     KeyExchangeMessage aliceKeyExchange = aliceSessionBuilder.process();
@@ -467,10 +467,10 @@ public class SessionBuilderTest extends TestCase {
   }
 
   public void testOptionalOneTimePreKey() throws Exception {
-    AxolotlStore   aliceStore          = new InMemoryAxolotlStore();
+    AxolotlStore   aliceStore          = new TestInMemoryAxolotlStore();
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
 
-    AxolotlStore bobStore = new InMemoryAxolotlStore();
+    AxolotlStore bobStore = new TestInMemoryAxolotlStore();
 
     ECKeyPair bobPreKeyPair            = Curve.generateKeyPair();
     ECKeyPair bobSignedPreKeyPair      = Curve.generateKeyPair();
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/SessionCipherTest.java b/tests/src/test/java/org/whispersystems/libaxolotl/SessionCipherTest.java
index a3b5db4c..41731982 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/SessionCipherTest.java
+++ b/tests/src/test/java/org/whispersystems/libaxolotl/SessionCipherTest.java
@@ -49,8 +49,8 @@ public class SessionCipherTest extends TestCase {
 
   private void runInteraction(SessionRecord aliceSessionRecord, SessionRecord bobSessionRecord)
       throws DuplicateMessageException, LegacyMessageException, InvalidMessageException, NoSuchAlgorithmException, NoSessionException {
-    AxolotlStore aliceStore = new InMemoryAxolotlStore();
-    AxolotlStore bobStore   = new InMemoryAxolotlStore();
+    AxolotlStore aliceStore = new TestInMemoryAxolotlStore();
+    AxolotlStore bobStore   = new TestInMemoryAxolotlStore();
 
     aliceStore.storeSession(new AxolotlAddress("+14159999999", 1), aliceSessionRecord);
     bobStore.storeSession(new AxolotlAddress("+14158888888", 1), bobSessionRecord);
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/SimultaneousInitiateTests.java b/tests/src/test/java/org/whispersystems/libaxolotl/SimultaneousInitiateTests.java
index 6e68508c..90e56d94 100644
--- a/tests/src/test/java/org/whispersystems/libaxolotl/SimultaneousInitiateTests.java
+++ b/tests/src/test/java/org/whispersystems/libaxolotl/SimultaneousInitiateTests.java
@@ -32,8 +32,8 @@ public class SimultaneousInitiateTests extends TestCase {
       InvalidMessageException, DuplicateMessageException, LegacyMessageException,
       InvalidKeyIdException, NoSessionException
   {
-    AxolotlStore aliceStore = new InMemoryAxolotlStore();
-    AxolotlStore bobStore   = new InMemoryAxolotlStore();
+    AxolotlStore aliceStore = new TestInMemoryAxolotlStore();
+    AxolotlStore bobStore   = new TestInMemoryAxolotlStore();
 
     PreKeyBundle alicePreKeyBundle = createAlicePreKeyBundle(aliceStore);
     PreKeyBundle bobPreKeyBundle = createBobPreKeyBundle(bobStore);
@@ -86,8 +86,8 @@ public class SimultaneousInitiateTests extends TestCase {
   }
 
   public void testLostSimultaneousInitiate() throws InvalidKeyException, UntrustedIdentityException, InvalidVersionException, InvalidMessageException, DuplicateMessageException, LegacyMessageException, InvalidKeyIdException, NoSessionException {
-    AxolotlStore aliceStore = new InMemoryAxolotlStore();
-    AxolotlStore bobStore   = new InMemoryAxolotlStore();
+    AxolotlStore aliceStore = new TestInMemoryAxolotlStore();
+    AxolotlStore bobStore   = new TestInMemoryAxolotlStore();
 
     PreKeyBundle alicePreKeyBundle = createAlicePreKeyBundle(aliceStore);
     PreKeyBundle bobPreKeyBundle = createBobPreKeyBundle(bobStore);
@@ -138,8 +138,8 @@ public class SimultaneousInitiateTests extends TestCase {
       InvalidMessageException, DuplicateMessageException, LegacyMessageException,
       InvalidKeyIdException, NoSessionException
   {
-    AxolotlStore aliceStore = new InMemoryAxolotlStore();
-    AxolotlStore bobStore   = new InMemoryAxolotlStore();
+    AxolotlStore aliceStore = new TestInMemoryAxolotlStore();
+    AxolotlStore bobStore   = new TestInMemoryAxolotlStore();
 
     PreKeyBundle alicePreKeyBundle = createAlicePreKeyBundle(aliceStore);
     PreKeyBundle bobPreKeyBundle = createBobPreKeyBundle(bobStore);
@@ -197,8 +197,8 @@ public class SimultaneousInitiateTests extends TestCase {
       InvalidMessageException, DuplicateMessageException, LegacyMessageException,
       InvalidKeyIdException, NoSessionException
   {
-    AxolotlStore aliceStore = new InMemoryAxolotlStore();
-    AxolotlStore bobStore   = new InMemoryAxolotlStore();
+    AxolotlStore aliceStore = new TestInMemoryAxolotlStore();
+    AxolotlStore bobStore   = new TestInMemoryAxolotlStore();
 
     PreKeyBundle alicePreKeyBundle = createAlicePreKeyBundle(aliceStore);
     PreKeyBundle bobPreKeyBundle = createBobPreKeyBundle(bobStore);
@@ -273,8 +273,8 @@ public class SimultaneousInitiateTests extends TestCase {
       InvalidMessageException, DuplicateMessageException, LegacyMessageException,
       InvalidKeyIdException, NoSessionException
   {
-    AxolotlStore aliceStore = new InMemoryAxolotlStore();
-    AxolotlStore bobStore   = new InMemoryAxolotlStore();
+    AxolotlStore aliceStore = new TestInMemoryAxolotlStore();
+    AxolotlStore bobStore   = new TestInMemoryAxolotlStore();
 
 
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
@@ -352,8 +352,8 @@ public class SimultaneousInitiateTests extends TestCase {
       InvalidMessageException, DuplicateMessageException, LegacyMessageException,
       InvalidKeyIdException, NoSessionException
   {
-    AxolotlStore aliceStore = new InMemoryAxolotlStore();
-    AxolotlStore bobStore   = new InMemoryAxolotlStore();
+    AxolotlStore aliceStore = new TestInMemoryAxolotlStore();
+    AxolotlStore bobStore   = new TestInMemoryAxolotlStore();
 
 
     SessionBuilder aliceSessionBuilder = new SessionBuilder(aliceStore, BOB_ADDRESS);
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/TestInMemoryAxolotlStore.java b/tests/src/test/java/org/whispersystems/libaxolotl/TestInMemoryAxolotlStore.java
new file mode 100644
index 00000000..8c4700fc
--- /dev/null
+++ b/tests/src/test/java/org/whispersystems/libaxolotl/TestInMemoryAxolotlStore.java
@@ -0,0 +1,22 @@
+package org.whispersystems.libaxolotl;
+
+import org.whispersystems.libaxolotl.ecc.Curve;
+import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+import org.whispersystems.libaxolotl.util.KeyHelper;
+
+public class TestInMemoryAxolotlStore extends org.whispersystems.libaxolotl.state.impl.InMemoryAxolotlStore {
+  public TestInMemoryAxolotlStore() {
+    super(generateIdentityKeyPair(), generateRegistrationId());
+  }
+
+  private static IdentityKeyPair generateIdentityKeyPair() {
+    ECKeyPair identityKeyPairKeys = Curve.generateKeyPair();
+
+    return new IdentityKeyPair(new IdentityKey(identityKeyPairKeys.getPublicKey()),
+                                               identityKeyPairKeys.getPrivateKey());
+  }
+
+  private static int generateRegistrationId() {
+    return KeyHelper.generateRegistrationId(false);
+  }
+}
diff --git a/tests/src/test/java/org/whispersystems/libaxolotl/TestInMemoryIdentityKeyStore.java b/tests/src/test/java/org/whispersystems/libaxolotl/TestInMemoryIdentityKeyStore.java
new file mode 100644
index 00000000..10a206d3
--- /dev/null
+++ b/tests/src/test/java/org/whispersystems/libaxolotl/TestInMemoryIdentityKeyStore.java
@@ -0,0 +1,23 @@
+package org.whispersystems.libaxolotl;
+
+import org.whispersystems.libaxolotl.ecc.Curve;
+import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+import org.whispersystems.libaxolotl.util.KeyHelper;
+
+public class TestInMemoryIdentityKeyStore extends org.whispersystems.libaxolotl.state.impl.InMemoryIdentityKeyStore {
+  public TestInMemoryIdentityKeyStore() {
+    super(generateIdentityKeyPair(), generateRegistrationId());
+  }
+
+  private static IdentityKeyPair generateIdentityKeyPair() {
+    ECKeyPair identityKeyPairKeys = Curve.generateKeyPair();
+
+    return new IdentityKeyPair(new IdentityKey(identityKeyPairKeys.getPublicKey()),
+                               identityKeyPairKeys.getPrivateKey());
+  }
+
+  private static int generateRegistrationId() {
+    return KeyHelper.generateRegistrationId(false);
+  }
+
+}