5 files changed, 25 insertions, 17 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;