Make invalidation tracker test friendly

This CL guards invalidation tracker from unwated failures which
can happen when database is closed.

Bug: 37160100
Test: InvalidationTrackerTest
Change-Id: I53511d89ab30f3dba01bac069a224c3fd5bc1ec0
diff --git a/room/integration-tests/testapp/src/androidTest/java/com/android/support/room/integration/testapp/migration/MigrationTest.java b/room/integration-tests/testapp/src/androidTest/java/com/android/support/room/integration/testapp/migration/MigrationTest.java
index ba77c85..8adfa3f 100644
--- a/room/integration-tests/testapp/src/androidTest/java/com/android/support/room/integration/testapp/migration/MigrationTest.java
+++ b/room/integration-tests/testapp/src/androidTest/java/com/android/support/room/integration/testapp/migration/MigrationTest.java
@@ -64,7 +64,7 @@
         db.close();
         MigrationDb migrationDb = getLatestDb();
         List<MigrationDb.Entity1> items = migrationDb.dao().loadAllEntity1s();
-        helper.closeWhenFinished(migrationDb.getDatabase());
+        helper.closeWhenFinished(migrationDb);
         assertThat(items.size(), is(1));
     }
 
@@ -97,7 +97,7 @@
         // trigger open
         db.beginTransaction();
         db.endTransaction();
-        helper.closeWhenFinished(db.getDatabase());
+        helper.closeWhenFinished(db);
         return db;
     }
 
diff --git a/room/runtime/src/main/java/com/android/support/room/InvalidationTracker.java b/room/runtime/src/main/java/com/android/support/room/InvalidationTracker.java
index 57e2f05..ccca026 100644
--- a/room/runtime/src/main/java/com/android/support/room/InvalidationTracker.java
+++ b/room/runtime/src/main/java/com/android/support/room/InvalidationTracker.java
@@ -17,6 +17,7 @@
 package com.android.support.room;
 
 import android.database.Cursor;
+import android.database.sqlite.SQLiteException;
 import android.support.annotation.NonNull;
 import android.support.annotation.Nullable;
 import android.support.annotation.RestrictTo;
@@ -99,7 +100,7 @@
 
     private final RoomDatabase mDatabase;
 
-    private AtomicBoolean mPendingRefresh = new AtomicBoolean(false);
+    AtomicBoolean mPendingRefresh = new AtomicBoolean(false);
 
     private volatile boolean mInitialized = false;
 
@@ -307,16 +308,16 @@
                     }
                     mObservedTableTracker.onSyncCompleted();
                 }
-            } catch (IllegalStateException exception) {
+            } catch (IllegalStateException | SQLiteException exception) {
                 // may happen if db is closed. just log.
-                Log.e(Room.LOG_TAG, "Cannot run invalidation tracker.", exception);
+                Log.e(Room.LOG_TAG, "Cannot run invalidation tracker. Is the db closed?",
+                        exception);
             }
         }
     };
 
     private boolean ensureInitialization() {
-        SupportSQLiteDatabase connection = mDatabase.getDatabase();
-        if (connection == null || !connection.isOpen()) {
+        if (!mDatabase.isOpen()) {
             return false;
         }
         if (!mInitialized) {
@@ -361,9 +362,10 @@
                 } finally {
                     cursor.close();
                 }
-            } catch (IllegalStateException exception) {
+            } catch (IllegalStateException | SQLiteException exception) {
                 // may happen if db is closed. just log.
-                Log.e(Room.LOG_TAG, "Cannot run invalidation tracker.", exception);
+                Log.e(Room.LOG_TAG, "Cannot run invalidation tracker. Is the db closed?",
+                        exception);
             }
             if (hasUpdatedTable) {
                 synchronized (mObserverMap) {
@@ -451,7 +453,7 @@
          * @param rest       More table names
          */
         @SuppressWarnings("unused")
-        public Observer(@NonNull String firstTable, String... rest) {
+        protected Observer(@NonNull String firstTable, String... rest) {
             mTables = Arrays.copyOf(rest, rest.length + 1);
             mTables[rest.length] = firstTable;
         }
diff --git a/room/runtime/src/main/java/com/android/support/room/RoomDatabase.java b/room/runtime/src/main/java/com/android/support/room/RoomDatabase.java
index 76e5b78..0ffada1 100644
--- a/room/runtime/src/main/java/com/android/support/room/RoomDatabase.java
+++ b/room/runtime/src/main/java/com/android/support/room/RoomDatabase.java
@@ -102,15 +102,22 @@
     protected abstract InvalidationTracker createInvalidationTracker();
 
     /**
-     * Returns the database connection. Note that, if the database is not opened yet, this method
-     * will return {code null}. You can use the {@link #getOpenHelper()} method to open the
-     * database.
+     * Returns true if database connection is open and initialized.
      *
-     * @return The database connection or {@code null} if it is not opened yet.
+     * @return true if the database connection is open, false otherwise.
      */
-    @Nullable
-    public SupportSQLiteDatabase getDatabase() {
-        return mDatabase;
+    public boolean isOpen() {
+        final SupportSQLiteDatabase db = mDatabase;
+        return db != null && db.isOpen();
+    }
+
+    /**
+     * Closes the database if it is already open.
+     */
+    public void close() {
+        if (isOpen()) {
+            mOpenHelper.close();
+        }
     }
 
     // Below, there are wrapper methods for SupportSQLiteDatabase. This helps us track which
diff --git a/room/runtime/src/test/java/com/android/support/room/InvalidationTrackerTest.java b/room/runtime/src/test/java/com/android/support/room/InvalidationTrackerTest.java
index ac70989..0e930a6 100644
--- a/room/runtime/src/test/java/com/android/support/room/InvalidationTrackerTest.java
+++ b/room/runtime/src/test/java/com/android/support/room/InvalidationTrackerTest.java
@@ -23,12 +23,14 @@
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.database.Cursor;
+import android.database.sqlite.SQLiteException;
 
 import com.android.support.apptoolkit.testing.JunitTaskExecutorRule;
 import com.android.support.db.SupportSQLiteDatabase;
@@ -54,6 +56,7 @@
 public class InvalidationTrackerTest {
     private InvalidationTracker mTracker;
     private RoomDatabase mRoomDatabase;
+    private SupportSQLiteOpenHelper mOpenHelper;
     @Rule
     public JunitTaskExecutorRule mTaskExecutorRule = new JunitTaskExecutorRule(1, true);
 
@@ -62,14 +65,13 @@
         mRoomDatabase = mock(RoomDatabase.class);
         SupportSQLiteDatabase sqliteDb = mock(SupportSQLiteDatabase.class);
         final SupportSQLiteStatement statement = mock(SupportSQLiteStatement.class);
-        SupportSQLiteOpenHelper openHelper = mock(SupportSQLiteOpenHelper.class);
+        mOpenHelper = mock(SupportSQLiteOpenHelper.class);
 
         doReturn(statement).when(sqliteDb).compileStatement(eq(InvalidationTracker.CLEANUP_SQL));
-        doReturn(sqliteDb).when(openHelper).getWritableDatabase();
-        doReturn(sqliteDb).when(mRoomDatabase).getDatabase();
-        doReturn(true).when(sqliteDb).isOpen();
+        doReturn(sqliteDb).when(mOpenHelper).getWritableDatabase();
+        doReturn(true).when(mRoomDatabase).isOpen();
         //noinspection ResultOfMethodCallIgnored
-        doReturn(openHelper).when(mRoomDatabase).getOpenHelper();
+        doReturn(mOpenHelper).when(mRoomDatabase).getOpenHelper();
 
         mTracker = new InvalidationTracker(mRoomDatabase, "a", "B", "i");
         mTracker.internalInit(sqliteDb);
@@ -199,6 +201,27 @@
         mTracker.addObserver(observer);
     }
 
+    @Test
+    public void closedDb() {
+        doThrow(new IllegalStateException("foo")).when(mOpenHelper).getWritableDatabase();
+        mTracker.addObserver(new LatchObserver(1, "a", "b"));
+        mTracker.syncTriggers();
+        mTracker.mRefreshRunnable.run();
+    }
+
+    @Test
+    public void closedDbAfterOpen() throws InterruptedException {
+        setVersions(3, 1);
+        mTracker.addObserver(new LatchObserver(1, "a", "b"));
+        mTracker.syncTriggers();
+        mTracker.mRefreshRunnable.run();
+        doThrow(new SQLiteException("foo")).when(mRoomDatabase).query(
+                Mockito.eq(InvalidationTracker.SELECT_UPDATED_TABLES_SQL),
+                any(String[].class));
+        mTracker.mPendingRefresh.set(true);
+        mTracker.mRefreshRunnable.run();
+    }
+
     /**
      * Key value pairs of VERSION, TABLE_ID
      */
diff --git a/room/testing/src/main/java/com/android/support/room/testing/MigrationTestHelper.java b/room/testing/src/main/java/com/android/support/room/testing/MigrationTestHelper.java
index 8c1d02d..8d4e4ab 100644
--- a/room/testing/src/main/java/com/android/support/room/testing/MigrationTestHelper.java
+++ b/room/testing/src/main/java/com/android/support/room/testing/MigrationTestHelper.java
@@ -77,6 +77,7 @@
     private final String mAssetsFolder;
     private final SupportSQLiteOpenHelper.Factory mOpenFactory;
     private List<WeakReference<SupportSQLiteDatabase>> mManagedDatabases = new ArrayList<>();
+    private List<WeakReference<RoomDatabase>> mManagedRoomDatabases = new ArrayList<>();
     private boolean mTestStarted;
 
     /**
@@ -199,6 +200,12 @@
                 }
             }
         }
+        for (WeakReference<RoomDatabase> dbRef : mManagedRoomDatabases) {
+            final RoomDatabase roomDatabase = dbRef.get();
+            if (roomDatabase != null) {
+                roomDatabase.close();
+            }
+        }
     }
 
     /**
@@ -218,6 +225,23 @@
         mManagedDatabases.add(new WeakReference<>(db));
     }
 
+    /**
+     * Registers a database connection to be automatically closed when the test finishes.
+     * <p>
+     * This only works if {@code MigrationTestHelper} is registered as a Junit test rule via
+     * {@link org.junit.Rule Rule} annotation.
+     *
+     * @param db The RoomDatabase instance which holds the database.
+     */
+    public void closeWhenFinished(RoomDatabase db) {
+        if (!mTestStarted) {
+            throw new IllegalStateException("You cannot register a database to be closed before"
+                    + " the test starts. Maybe you forgot to annotate MigrationTestHelper as a"
+                    + " test rule? (@Rule)");
+        }
+        mManagedRoomDatabases.add(new WeakReference<>(db));
+    }
+
     private SchemaBundle loadSchema(int version) throws IOException {
         InputStream input = mContext.getAssets().open(mAssetsFolder + "/" + version + ".json");
         return SchemaBundle.deserialize(input);