diff --git a/constants.gradle b/constants.gradle
index f6fc3d3d..0cfd59f3 100644
--- a/constants.gradle
+++ b/constants.gradle
@@ -16,7 +16,7 @@
 // TODO migrate to version catalogs
 ext {
     groupId = "com.grab.grazel"
-    versionName = "0.3.0"
+    versionName = "0.3.1"
 
     kotlinVersion = "1.4.31"
     agpVersion = "4.1.3"
@@ -36,4 +36,4 @@ ext {
     dokkaVersion = "1.4.32"
 
     website = "https://grab.github.io/Grazel/"
-}
\ No newline at end of file
+}
diff --git a/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/KotlinUnitTestDataExtractor.kt b/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/KotlinUnitTestDataExtractor.kt
index d2ea2c90..6edbe3f2 100644
--- a/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/KotlinUnitTestDataExtractor.kt
+++ b/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/KotlinUnitTestDataExtractor.kt
@@ -59,7 +59,6 @@ internal class DefaultKotlinUnitTestDataExtractor @Inject constructor(
         val deps: List<BazelDependency> = buildList {
             addAll(projectDependencyGraphs.directProjectDependencies(project, ConfigurationScope.TEST))
             addAll(dependenciesDataSource.collectMavenDeps(project, ConfigurationScope.TEST))
-            addAll(project.androidJarDeps())
             addAll(project.kotlinParcelizeDeps())
             if (projectDependency.toString() != associate.toString()) {
                 add(projectDependency)
@@ -69,7 +68,8 @@ internal class DefaultKotlinUnitTestDataExtractor @Inject constructor(
             name = FORMAT_UNIT_TEST_NAME.format(project.name),
             srcs = srcs,
             deps = deps,
-            associates = buildList { associate?.let(::add) }
+            associates = buildList { associate?.let(::add) },
+            hasAndroidJarDep = project.hasAndroidJarDep(),
         )
     }
 
@@ -84,14 +84,9 @@ internal class DefaultKotlinUnitTestDataExtractor @Inject constructor(
     }
 }
 
-internal fun Project.androidJarDeps(): List<BazelDependency> {
-    return if (configurations.findByName("compileOnly")
-            ?.dependencies
-            ?.filterIsInstance<DefaultSelfResolvingDependency>()
-            ?.any { dep -> dep.files.any { it.name.contains("android.jar") } } == true
-    ) {
-        listOf(BazelDependency.StringDependency("//shared_versions:android_sdk"))
-    } else {
-        emptyList()
-    }
-}
\ No newline at end of file
+internal fun Project.hasAndroidJarDep(): Boolean {
+    return configurations.findByName("compileOnly")
+        ?.dependencies
+        ?.filterIsInstance<DefaultSelfResolvingDependency>()
+        ?.any {  dep -> dep.files.any { it.name.contains("android.jar") } } == true
+}
diff --git a/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/UnitTestData.kt b/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/UnitTestData.kt
index 37e52398..d38b6afe 100644
--- a/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/UnitTestData.kt
+++ b/grazel-gradle-plugin/src/main/kotlin/com/grab/grazel/migrate/kotlin/UnitTestData.kt
@@ -1,19 +1,31 @@
 package com.grab.grazel.migrate.kotlin
 
 import com.grab.grazel.bazel.starlark.BazelDependency
+import com.grab.grazel.migrate.BazelBuildTarget
+import com.grab.grazel.migrate.android.AndroidUnitTestTarget
 
 data class UnitTestData(
     val name: String,
     val srcs: List<String>,
     val deps: List<BazelDependency>,
-    val associates: List<BazelDependency>
+    val associates: List<BazelDependency>,
+    val hasAndroidJarDep: Boolean = false,
 )
 
-internal fun UnitTestData.toUnitTestTarget(): UnitTestTarget {
-    return UnitTestTarget(
-        name = name,
-        srcs = srcs,
-        deps = deps,
-        associates = associates
-    )
-}
+internal fun UnitTestData.toUnitTestTarget(): BazelBuildTarget =
+    if (hasAndroidJarDep) {
+        AndroidUnitTestTarget(
+            name = name,
+            srcs = srcs,
+            deps = deps,
+            associates = associates,
+            customPackage = "",
+        )
+    } else {
+        UnitTestTarget(
+            name = name,
+            srcs = srcs,
+            deps = deps,
+            associates = associates
+        )
+    }