Support robolectric tests for multi process datastore

MultiProcessDataStore uses native code, which does not work well
on robolectric (and the class is not even available on the JVM
target). This CL updates it to bypass loading the native library
if we cannot load it AND we are not on dalvik.

I've created a new sourceset to combine MultiProcess tests so that
they are run on robolectric AND android device to ensure we don't
break this in the future.

I've also added proguard rules to ensure r8 can remove this fallback
implementation when compiling for devices.

Test: PROJECT_PREFIX=:datastore ./gradlew :datastore:datastore-core:cleanTestReleaseUnitTest :datastore:datastore-core:testReleaseUnitTest :datastore:datastore-core:cC
Fixes: b/352047731
Change-Id: I6d1f088b3a237e04914215d4864421d4105d7a7b
diff --git a/datastore/datastore-core/build.gradle b/datastore/datastore-core/build.gradle
index ae103b5..a722f28 100644
--- a/datastore/datastore-core/build.gradle
+++ b/datastore/datastore-core/build.gradle
@@ -39,6 +39,9 @@
         }
     }
     namespace = "androidx.datastore.core"
+    buildTypes.configureEach {
+        consumerProguardFiles("proguard-rules.pro")
+    }
 }
 
 androidXMultiplatform {
@@ -101,6 +104,11 @@
                 implementation(libs.kotlinTestJunit)
             }
         }
+        // tests that are run on AndroidUnit and instrumented targets.
+        // similar to jvmAndroidTest, except, they don't run on JVM outside Robolectric.
+        androidUnitAndInstrumentedTest {
+            dependsOn(jvmAndroidTest)
+        }
 
         jvmTest {
             dependsOn(jvmAndroidTest)
@@ -108,19 +116,26 @@
 
         androidUnitTest {
             dependsOn(jvmAndroidTest)
+            dependsOn(androidUnitAndInstrumentedTest)
             dependencies {
                 implementation(libs.protobufLite)
+                implementation(libs.testRules)
+                implementation(libs.testRunner)
+                implementation(libs.robolectric)
+                implementation(libs.testExtJunit)
             }
         }
 
         androidInstrumentedTest {
             dependsOn(jvmAndroidTest)
+            dependsOn(androidUnitAndInstrumentedTest)
             dependencies {
                 implementation(libs.truth)
                 implementation(project(":internal-testutils-truth"))
                 implementation(libs.testRunner)
                 implementation(libs.testCore)
                 implementation("androidx.lifecycle:lifecycle-service:2.6.1")
+                implementation(libs.testExtJunit)
 
                 // Workaround bug in 1.8.0, was supposed be fixed in RC2/final, but apparently not.
                 implementation(libs.kotlinTestJunit)
diff --git a/datastore/datastore-core/proguard-rules.pro b/datastore/datastore-core/proguard-rules.pro
new file mode 100644
index 0000000..f43b709
--- /dev/null
+++ b/datastore/datastore-core/proguard-rules.pro
@@ -0,0 +1,5 @@
+# make sure r8 can remove the fake counter that is only in the code to handle the
+# robolectric tests.
+-assumenosideeffects class androidx.datastore.core.SharedCounter$Factory {
+    private boolean isDalvik() return true;
+}
diff --git a/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/MultiProcessCoordinator.android.kt b/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/MultiProcessCoordinator.android.kt
index c753218..3490ad0 100644
--- a/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/MultiProcessCoordinator.android.kt
+++ b/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/MultiProcessCoordinator.android.kt
@@ -123,7 +123,6 @@
     }
 
     private val lazySharedCounter = lazy {
-        SharedCounter.loadLib()
         SharedCounter.create {
             val versionFile = fileWithSuffix(VERSION_SUFFIX)
             versionFile.createIfNotExists()
diff --git a/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/SharedCounter.android.kt b/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/SharedCounter.android.kt
index f81e8c3..d415060 100644
--- a/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/SharedCounter.android.kt
+++ b/datastore/datastore-core/src/androidMain/kotlin/androidx/datastore/core/SharedCounter.android.kt
@@ -19,6 +19,7 @@
 import android.os.ParcelFileDescriptor
 import java.io.File
 import java.io.IOException
+import java.util.concurrent.atomic.AtomicInteger
 
 /** Put the JNI methods in a separate class to make them internal to the package. */
 internal class NativeSharedCounter {
@@ -37,26 +38,77 @@
  * of the `datastore-multiprocess` AAR artifact, users don't need extra steps other than adding it
  * as dependency.
  */
-internal class SharedCounter
-private constructor(
-    /** The memory address to be mapped. */
-    private val mappedAddress: Long
-) {
+internal interface SharedCounter {
+    fun getValue(): Int
 
-    fun getValue(): Int {
-        return nativeSharedCounter.nativeGetCounterValue(mappedAddress)
+    fun incrementAndGetValue(): Int
+
+    private class RealSharedCounter(
+        private val nativeSharedCounter: NativeSharedCounter,
+        /** The memory address to be mapped. */
+        private val mappedAddress: Long
+    ) : SharedCounter {
+        override fun getValue(): Int {
+            return nativeSharedCounter.nativeGetCounterValue(mappedAddress)
+        }
+
+        override fun incrementAndGetValue(): Int {
+            return nativeSharedCounter.nativeIncrementAndGetCounterValue(mappedAddress)
+        }
     }
 
-    fun incrementAndGetValue(): Int {
-        return nativeSharedCounter.nativeIncrementAndGetCounterValue(mappedAddress)
+    /** Shared counter implementation that is used when running Robolectric tests. */
+    private class ShadowSharedCounter : SharedCounter {
+        private val value = AtomicInteger(0)
+
+        override fun getValue(): Int {
+            return value.get()
+        }
+
+        override fun incrementAndGetValue(): Int {
+            return value.incrementAndGet()
+        }
     }
 
     companion object Factory {
-        internal val nativeSharedCounter = NativeSharedCounter()
-
-        fun loadLib() = System.loadLibrary("datastore_shared_counter")
+        private val nativeSharedCounter: NativeSharedCounter? =
+            try {
+                System.loadLibrary("datastore_shared_counter")
+                NativeSharedCounter()
+            } catch (th: Throwable) {
+                /**
+                 * Currently this native library is only available for Android, it should not be
+                 * loaded on host platforms, e.g. Robolectric.
+                 */
+                if (isDalvik()) {
+                    // we should always be able to load it on dalvik
+                    throw th
+                } else {
+                    // probably running on robolectric, ignore.
+                    null
+                }
+            }
 
         private fun createCounterFromFd(pfd: ParcelFileDescriptor): SharedCounter {
+            if (nativeSharedCounter == null) {
+                // don't remove the following isDalvik check, it helps r8 cleanup the
+                // ShadowSharedCounter code for android.
+                if (!isDalvik()) {
+                    // if it is null, we are not on Android so just use an in
+                    // process shared counter as multi-process is not testable on
+                    // Robolectric.
+                    return ShadowSharedCounter()
+                }
+                // This actually will enver happen, because class creation would throw when
+                // initializing nativeSharedCounter. But having it here helps future proofing as
+                // well as the static analyzer.
+                error(
+                    """
+                    DataStore failed to load the native library to create SharedCounter.
+                """
+                        .trimIndent()
+                )
+            }
             val nativeFd = pfd.getFd()
             if (nativeSharedCounter.nativeTruncateFile(nativeFd) != 0) {
                 throw IOException("Failed to truncate counter file")
@@ -65,7 +117,7 @@
             if (address < 0) {
                 throw IOException("Failed to mmap counter file")
             }
-            return SharedCounter(address)
+            return RealSharedCounter(nativeSharedCounter, address)
         }
 
         internal fun create(produceFile: () -> File): SharedCounter {
@@ -82,5 +134,10 @@
                 pfd?.close()
             }
         }
+
+        /** If you change this method, make sure to update the proguard rule. */
+        private fun isDalvik(): Boolean {
+            return "dalvik".equals(System.getProperty("java.vm.name"), ignoreCase = true)
+        }
     }
 }
diff --git a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreFactoryTest.kt b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreFactoryTest.kt
similarity index 96%
rename from datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreFactoryTest.kt
rename to datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreFactoryTest.kt
index 1c760e4..7e5a01f 100644
--- a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreFactoryTest.kt
+++ b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreFactoryTest.kt
@@ -18,7 +18,8 @@
 
 import androidx.datastore.TestingSerializerConfig
 import androidx.datastore.core.handlers.ReplaceFileCorruptionHandler
-import com.google.common.truth.Truth.assertThat
+import androidx.kruth.assertThat
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import java.io.File
 import java.util.concurrent.TimeUnit
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -32,11 +33,10 @@
 import org.junit.rules.TemporaryFolder
 import org.junit.rules.Timeout
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 /** A duplicated test of {@link androidx.datastore.core.DataStoreFactoryTest} with minor changes. */
 @ExperimentalCoroutinesApi
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class MultiProcessDataStoreFactoryTest {
     @get:Rule val timeout = Timeout(10, TimeUnit.SECONDS)
 
diff --git a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessFileTest.kt b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessFileTest.kt
similarity index 93%
rename from datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessFileTest.kt
rename to datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessFileTest.kt
index 7c95a65..ddcec6d8 100644
--- a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessFileTest.kt
+++ b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessFileTest.kt
@@ -18,8 +18,9 @@
 
 import androidx.datastore.FileTestIO
 import androidx.datastore.JavaIOFile
-import androidx.testutils.assertThrows
-import com.google.common.truth.Truth
+import androidx.kruth.assertThat
+import androidx.kruth.assertThrows
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import java.io.File
 import java.io.IOException
 import java.io.InputStream
@@ -39,11 +40,14 @@
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
+import org.junit.runner.RunWith
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @InternalCoroutinesApi
+@RunWith(AndroidJUnit4::class)
 class MultiProcessDataStoreSingleProcessFileTest :
     MultiProcessDataStoreSingleProcessTest<JavaIOFile>(FileTestIO()) {
+
     override fun getJavaFile(file: JavaIOFile): File {
         return file.file
     }
@@ -55,8 +59,8 @@
         testFile.file.setReadable(false)
         val result = runCatching { store.data.first() }
 
-        Truth.assertThat(result.exceptionOrNull()).isInstanceOf(IOException::class.java)
-        Truth.assertThat(result.exceptionOrNull())
+        assertThat(result.exceptionOrNull()).isInstanceOf<IOException>()
+        assertThat(result.exceptionOrNull())
             .hasCauseThat()
             .hasMessageThat()
             .contains("Permission denied")
@@ -74,7 +78,7 @@
             .contains("Permission denied")
 
         testFile.file.setReadable(true)
-        Truth.assertThat(store.data.first()).isEqualTo(0)
+        assertThat(store.data.first()).isEqualTo(0)
     }
 
     @Test
diff --git a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessOkioTest.kt b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessOkioTest.kt
similarity index 91%
rename from datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessOkioTest.kt
rename to datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessOkioTest.kt
index 8dbe972..4f7e955 100644
--- a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessOkioTest.kt
+++ b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessOkioTest.kt
@@ -18,14 +18,17 @@
 
 import androidx.datastore.OkioPath
 import androidx.datastore.OkioTestIO
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import java.io.File
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.ObsoleteCoroutinesApi
+import org.junit.runner.RunWith
 
 @OptIn(ExperimentalCoroutinesApi::class, ObsoleteCoroutinesApi::class, FlowPreview::class)
 @InternalCoroutinesApi
+@RunWith(AndroidJUnit4::class)
 class MultiProcessDataStoreSingleProcessOkioTest :
     MultiProcessDataStoreSingleProcessTest<OkioPath>(OkioTestIO()) {
     override fun getJavaFile(file: OkioPath): File {
diff --git a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessTest.kt b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessTest.kt
similarity index 99%
rename from datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessTest.kt
rename to datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessTest.kt
index 6065045..bd414b6 100644
--- a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessTest.kt
+++ b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/MultiProcessDataStoreSingleProcessTest.kt
@@ -21,10 +21,11 @@
 import androidx.datastore.TestIO
 import androidx.datastore.TestingSerializerConfig
 import androidx.datastore.core.handlers.NoOpCorruptionHandler
+import androidx.kruth.assertThat
+import androidx.kruth.assertThrows
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.LargeTest
-import androidx.testutils.assertThrows
-import com.google.common.truth.Truth.assertThat
 import java.io.File
 import java.io.InputStream
 import java.io.OutputStream
@@ -61,7 +62,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 /**
  * A testing class based on duplicate from "SingleProcessDataStoreTest" that only tests the features
@@ -70,7 +70,7 @@
 @OptIn(DelicateCoroutinesApi::class)
 @ExperimentalCoroutinesApi
 @LargeTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 abstract class MultiProcessDataStoreSingleProcessTest<F : TestFile<F>>(
     protected val testIO: TestIO<F, *>
 ) {
diff --git a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/SharedCounterTest.kt b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/SharedCounterTest.kt
similarity index 89%
rename from datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/SharedCounterTest.kt
rename to datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/SharedCounterTest.kt
index c02ca7a..9363626 100644
--- a/datastore/datastore-core/src/androidInstrumentedTest/kotlin/androidx/datastore/core/SharedCounterTest.kt
+++ b/datastore/datastore-core/src/androidUnitAndInstrumentedTest/kotlin/androidx/datastore/core/SharedCounterTest.kt
@@ -16,12 +16,11 @@
 
 package androidx.datastore.core
 
+import androidx.kruth.assertThat
+import androidx.kruth.assertThrows
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
-import androidx.testutils.assertThrows
-import com.google.common.truth.Truth.assertThat
 import java.io.File
-import java.io.IOException
-import kotlin.collections.MutableSet
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.async
 import kotlinx.coroutines.launch
@@ -31,19 +30,11 @@
 import org.junit.Test
 import org.junit.rules.TemporaryFolder
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @ExperimentalCoroutinesApi
 @MediumTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class SharedCounterTest {
-
-    companion object {
-        init {
-            SharedCounter.loadLib()
-        }
-    }
-
     @get:Rule val tempFolder = TemporaryFolder()
     private lateinit var testFile: File
 
@@ -62,7 +53,7 @@
     fun testCreate_failure() {
         val tempFile = tempFolder.newFile()
         tempFile.setReadable(false)
-        assertThrows(IOException::class.java) { SharedCounter.create { tempFile } }
+        assertThrows<IOException> { SharedCounter.create { tempFile } }
     }
 
     @Test
diff --git a/datastore/datastore-sampleapp/build.gradle b/datastore/datastore-sampleapp/build.gradle
index b3e5bc0..66cba3e 100644
--- a/datastore/datastore-sampleapp/build.gradle
+++ b/datastore/datastore-sampleapp/build.gradle
@@ -55,5 +55,11 @@
 
 android {
     namespace = "com.example.datastoresampleapp"
+    buildTypes {
+        release {
+            minifyEnabled true
+            proguardFiles getDefaultProguardFile('proguard-android-optimize.txt'), 'proguard-rules.pro'
+        }
+    }
 }
 
diff --git a/datastore/datastore-sampleapp/proguard-rules.pro b/datastore/datastore-sampleapp/proguard-rules.pro
new file mode 100644
index 0000000..a4c0298
--- /dev/null
+++ b/datastore/datastore-sampleapp/proguard-rules.pro
@@ -0,0 +1 @@
+-keepnames class ** { *; }
diff --git a/datastore/datastore-sampleapp/src/main/java/com/example/datastoresampleapp/SettingsFragment.kt b/datastore/datastore-sampleapp/src/main/java/com/example/datastoresampleapp/SettingsFragment.kt
index 69a0e92..268d542 100644
--- a/datastore/datastore-sampleapp/src/main/java/com/example/datastoresampleapp/SettingsFragment.kt
+++ b/datastore/datastore-sampleapp/src/main/java/com/example/datastoresampleapp/SettingsFragment.kt
@@ -23,7 +23,7 @@
 import androidx.appcompat.app.AppCompatActivity
 import androidx.datastore.core.CorruptionException
 import androidx.datastore.core.DataStore
-import androidx.datastore.core.DataStoreFactory
+import androidx.datastore.core.MultiProcessDataStoreFactory
 import androidx.datastore.core.Serializer
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.lifecycleScope
@@ -72,7 +72,7 @@
     private val PROTO_STORE_FILE_NAME = "datastore_test_app.pb"
 
     private val settingsStore: DataStore<Settings> by lazy {
-        DataStoreFactory.create(serializer = SettingsSerializer) {
+        MultiProcessDataStoreFactory.create(serializer = SettingsSerializer) {
             File(requireActivity().applicationContext.filesDir, PROTO_STORE_FILE_NAME)
         }
     }