diff --git a/library-ktx/src/test/java/me/drakeet/multitype/MultiTypeTest.kt b/library-ktx/src/test/java/me/drakeet/multitype/MultiTypeTest.kt
index 4506fdee..95bb7167 100644
--- a/library-ktx/src/test/java/me/drakeet/multitype/MultiTypeTest.kt
+++ b/library-ktx/src/test/java/me/drakeet/multitype/MultiTypeTest.kt
@@ -16,7 +16,8 @@
 
 package me.drakeet.multitype
 
-import junit.framework.Assert
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
@@ -34,14 +35,14 @@ class MultiTypeTest {
   @Test
   fun shouldEqualToRegisteredKClass() {
     adapter.register(String::class, StringViewBinder())
-    Assert.assertEquals(adapter.typePool.getClass(0), String::class.java)
+    assertEquals(adapter.typePool.getClass(0), String::class.java)
   }
 
 
   @Test
   fun shouldEqualToRegisteredKClass_Reified() {
     adapter.register(StringViewBinder())
-    Assert.assertEquals(adapter.typePool.getClass(0), String::class.java)
+    assertEquals(adapter.typePool.getClass(0), String::class.java)
   }
 
 
@@ -50,27 +51,27 @@ class MultiTypeTest {
     adapter.register(String::class)
         .to(StringViewBinder())
         .withLinker(simpleLinker)
-    Assert.assertEquals(adapter.typePool.getClass(0), String::class.java)
+    assertEquals(adapter.typePool.getClass(0), String::class.java)
   }
 
 
   @Test
   fun shouldEqualToRegisteredKClass_TypePool() {
     adapter.typePool.register(String::class, StringViewBinder(), simpleLinker)
-    Assert.assertEquals(adapter.typePool.getClass(0), String::class.java)
+    assertEquals(adapter.typePool.getClass(0), String::class.java)
   }
 
 
   @Test
   fun shouldUnregisterKClass_TypePool() {
     adapter.typePool.register(String::class, StringViewBinder(), simpleLinker)
-    Assert.assertTrue(adapter.typePool.unregister(String::class))
+    assertTrue(adapter.typePool.unregister(String::class))
   }
 
 
   @Test
   fun shouldEqualToRegisteredFirstKClass_TypePool() {
     adapter.typePool.register(String::class, StringViewBinder(), simpleLinker)
-    Assert.assertEquals(adapter.typePool.firstIndexOf(String::class), 0)
+    assertEquals(adapter.typePool.firstIndexOf(String::class), 0)
   }
 }
diff --git a/library/src/test/java/me/drakeet/multitype/ItemViewBinderTest.java b/library/src/test/java/me/drakeet/multitype/ItemViewBinderTest.java
index f3a647bf..56265b3d 100644
--- a/library/src/test/java/me/drakeet/multitype/ItemViewBinderTest.java
+++ b/library/src/test/java/me/drakeet/multitype/ItemViewBinderTest.java
@@ -63,7 +63,6 @@ public void shouldThrowIllegalStateException() {
 
 
   private static class TestItemViewBinder extends me.drakeet.multitype.TestItemViewBinder {
-
     private void notifyTestItemAdded() {
       assertNotNull(getAdapter().toString());
     }
diff --git a/sample/build.gradle b/sample/build.gradle
index 02fb8e1d..0d98afe8 100644
--- a/sample/build.gradle
+++ b/sample/build.gradle
@@ -45,6 +45,7 @@ dependencies {
   implementation project(':library')
   implementation project(':library-ktx')
   androidTestImplementation 'androidx.test.espresso:espresso-core:3.1.0-alpha4'
+  androidTestImplementation 'androidx.test:rules:1.1.0-alpha4'
   androidTestImplementation "androidx.annotation:annotation:${gradle.androidXVersion}"
   implementation "androidx.appcompat:appcompat:${gradle.androidXVersion}"
   implementation "androidx.recyclerview:recyclerview:${gradle.androidXVersion}"