Use NoteDb consistently in query tests

InMemoryModule was missing the binding to use
NotesMigrationSchemaFactory, which is responsible for wrapping the
ReviewDb based on NotesMigration state. Additionally,
AbstractQueryChangesTest was keeping a reference to the underlying
InMemoryDatabase instance, where it should be opening from the injected
SchemaFactory.

Change-Id: I28a2bf40623fc306f9bbf21808cea609c561b5e7

gerrit-server/src/test/java/com/google/gerrit/server/query/change/AbstractQueryChangesTest.java

@@ -97,6 +97,7 @@ import com.google.gerrit.testutil.InMemoryDatabase;
 import com.google.gerrit.testutil.InMemoryRepositoryManager;
 import com.google.gerrit.testutil.InMemoryRepositoryManager.Repo;
 import com.google.gerrit.testutil.TestTimeUtil;
+import com.google.gwtorm.server.SchemaFactory;
 import com.google.inject.Inject;
 import com.google.inject.Injector;
 import com.google.inject.Provider;
@@ -143,7 +144,6 @@ public abstract class AbstractQueryChangesTest extends GerritServerTests {
   @Inject protected ChangeIndexCollection indexes;
   @Inject protected ChangeIndexer indexer;
   @Inject protected IndexConfig indexConfig;
-  @Inject protected InMemoryDatabase schemaFactory;
   @Inject protected InMemoryRepositoryManager repoManager;
   @Inject protected InternalChangeQuery internalChangeQuery;
   @Inject protected ChangeNotes.Factory notesFactory;
@@ -152,11 +152,15 @@ public abstract class AbstractQueryChangesTest extends GerritServerTests {
   @Inject protected ChangeControl.GenericFactory changeControlFactory;
   @Inject protected ChangeQueryProcessor queryProcessor;
   @Inject protected SchemaCreator schemaCreator;
+  @Inject protected SchemaFactory<ReviewDb> schemaFactory;
   @Inject protected Sequences seq;
   @Inject protected ThreadLocalRequestContext requestContext;
   @Inject protected ProjectCache projectCache;
   @Inject protected MetaDataUpdate.Server metaDataUpdateFactory;
 
+  // Only for use in setting up/tearing down injector; other users should use schemaFactory.
+  @Inject private InMemoryDatabase inMemoryDatabase;
+
   protected Injector injector;
   protected LifecycleManager lifecycle;
   protected ReviewDb db;
@@ -178,8 +182,10 @@ public abstract class AbstractQueryChangesTest extends GerritServerTests {
   }
 
   protected void setUpDatabase() throws Exception {
+    try (ReviewDb underlyingDb = inMemoryDatabase.getDatabase().open()) {
+      schemaCreator.create(underlyingDb);
+    }
     db = schemaFactory.open();
-    schemaCreator.create(db);
 
     userId = accountManager.authenticate(AuthRequest.forUser("user")).getAccountId();
     Account userAccount = db.accounts().get(userId);
@@ -213,7 +219,7 @@ public abstract class AbstractQueryChangesTest extends GerritServerTests {
     if (db != null) {
       db.close();
     }
-    InMemoryDatabase.drop(schemaFactory);
+    InMemoryDatabase.drop(inMemoryDatabase);
   }
 
   @Before
@@ -1370,7 +1376,8 @@ public abstract class AbstractQueryChangesTest extends GerritServerTests {
     allUsers.update(draftsRef.getName(), draftsRef.getObjectId());
     assertThat(allUsers.getRepository().exactRef(draftsRef.getName())).isNotNull();
 
-    if (PrimaryStorage.of(change) == PrimaryStorage.REVIEW_DB) {
+    if (PrimaryStorage.of(change) == PrimaryStorage.REVIEW_DB
+        && !notesMigration.disableChangeReviewDb()) {
       // Record draft ref in noteDbState as well.
       ReviewDb db = ReviewDbUtil.unwrapDb(this.db);
       change = db.changes().get(id);

gerrit-server/src/test/java/com/google/gerrit/testutil/InMemoryModule.java

@@ -60,6 +60,8 @@ import com.google.gerrit.server.notedb.NotesMigration;
 import com.google.gerrit.server.patch.DiffExecutor;
 import com.google.gerrit.server.schema.DataSourceType;
 import com.google.gerrit.server.schema.H2AccountPatchReviewStore;
+import com.google.gerrit.server.schema.NotesMigrationSchemaFactory;
+import com.google.gerrit.server.schema.ReviewDbFactory;
 import com.google.gerrit.server.schema.SchemaCreator;
 import com.google.gerrit.server.securestore.DefaultSecureStore;
 import com.google.gerrit.server.securestore.SecureStore;
@@ -70,6 +72,7 @@ import com.google.gwtorm.server.SchemaFactory;
 import com.google.inject.AbstractModule;
 import com.google.inject.Guice;
 import com.google.inject.Injector;
+import com.google.inject.Key;
 import com.google.inject.Module;
 import com.google.inject.Provider;
 import com.google.inject.Provides;
@@ -172,13 +175,15 @@ public class InMemoryModule extends FactoryModule {
     bind(ListeningExecutorService.class)
         .annotatedWith(ChangeUpdateExecutor.class)
         .toInstance(MoreExecutors.newDirectExecutorService());
-
     bind(DataSourceType.class).to(InMemoryH2Type.class);
-    bind(new TypeLiteral<SchemaFactory<ReviewDb>>() {}).to(InMemoryDatabase.class);
     bind(ChangeBundleReader.class).to(GwtormChangeBundleReader.class);
-
     bind(SecureStore.class).to(DefaultSecureStore.class);
 
+    TypeLiteral<SchemaFactory<ReviewDb>> schemaFactory =
+        new TypeLiteral<SchemaFactory<ReviewDb>>() {};
+    bind(schemaFactory).to(NotesMigrationSchemaFactory.class);
+    bind(Key.get(schemaFactory, ReviewDbFactory.class)).to(InMemoryDatabase.class);
+
     install(NoSshKeyCache.module());
     install(
         new CanonicalWebUrlModule() {