diff --git a/library/self.gradle b/library/self.gradle
index 6a405e2..69f9654 100644
--- a/library/self.gradle
+++ b/library/self.gradle
@@ -38,8 +38,8 @@ android.sourceSets {
         java.srcDirs = ['test/robotest/src']
 
         dependencies {
-            testCompile 'org.robolectric:robolectric:3.+'
-            testCompile 'org.robolectric:shadows-core:3.+'
+            testCompile 'org.robolectric:robolectric:3.4.2'
+            testCompile 'org.robolectric:framework:3.4.2'
             testCompile 'junit:junit:4.+'
             testCompile 'org.mockito:mockito-core:1.9.5'
             // Workaround for https://github.com/robolectric/robolectric/issues/2566
diff --git a/library/test/robotest/src/com/android/setupwizardlib/util/PartnerTest.java b/library/test/robotest/src/com/android/setupwizardlib/util/PartnerTest.java
index f47eef1..aeb678f 100644
--- a/library/test/robotest/src/com/android/setupwizardlib/util/PartnerTest.java
+++ b/library/test/robotest/src/com/android/setupwizardlib/util/PartnerTest.java
@@ -31,6 +31,7 @@ import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.ResolveInfo;
 import android.content.res.Resources;
 import android.os.Build.VERSION;
@@ -40,20 +41,25 @@ import com.android.setupwizardlib.BuildConfig;
 import com.android.setupwizardlib.R;
 import com.android.setupwizardlib.robolectric.SuwLibRobolectricTestRunner;
 import com.android.setupwizardlib.util.Partner.ResourceEntry;
+import com.android.setupwizardlib.util.PartnerTest.ShadowApplicationPackageManager;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RuntimeEnvironment;
+import org.robolectric.Shadows;
 import org.robolectric.annotation.Config;
-import org.robolectric.res.builder.DefaultPackageManager;
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
 import org.robolectric.shadows.ShadowResources;
 
 import java.util.Arrays;
 import java.util.Collections;
 
 @RunWith(SuwLibRobolectricTestRunner.class)
-@Config(constants = BuildConfig.class, sdk = { Config.OLDEST_SDK, Config.NEWEST_SDK })
+@Config(
+        constants = BuildConfig.class,
+        sdk = { Config.OLDEST_SDK, Config.NEWEST_SDK },
+        shadows = ShadowApplicationPackageManager.class)
 public class PartnerTest {
 
     private static final String ACTION_PARTNER_CUSTOMIZATION =
@@ -62,7 +68,7 @@ public class PartnerTest {
     private Context mContext;
     private Resources mPartnerResources;
 
-    private TestPackageManager mPackageManager;
+    private ShadowApplicationPackageManager mPackageManager;
 
     @Before
     public void setUp() throws Exception {
@@ -71,8 +77,9 @@ public class PartnerTest {
         mContext = spy(application);
         mPartnerResources = spy(ShadowResources.getSystem());
 
-        mPackageManager = new TestPackageManager();
-        RuntimeEnvironment.setRobolectricPackageManager(mPackageManager);
+        mPackageManager =
+                (ShadowApplicationPackageManager) Shadows.shadowOf(application.getPackageManager());
+        mPackageManager.partnerResources = mPartnerResources;
     }
 
     @Test
@@ -173,13 +180,18 @@ public class PartnerTest {
         return info;
     }
 
-    private class TestPackageManager extends DefaultPackageManager {
+    @Implements(className = "android.app.ApplicationPackageManager")
+    public static class ShadowApplicationPackageManager extends
+            org.robolectric.shadows.ShadowApplicationPackageManager {
 
+        public Resources partnerResources;
+
+        @Implementation
         @Override
         public Resources getResourcesForApplication(ApplicationInfo app)
                 throws NameNotFoundException {
             if (app != null && "test.partner.package".equals(app.packageName)) {
-                return mPartnerResources;
+                return partnerResources;
             } else {
                 return super.getResourcesForApplication(app);
             }
diff --git a/library/test/robotest/src/com/android/setupwizardlib/view/IllustrationVideoViewTest.java b/library/test/robotest/src/com/android/setupwizardlib/view/IllustrationVideoViewTest.java
index ffa228d..ddf59ca 100644
--- a/library/test/robotest/src/com/android/setupwizardlib/view/IllustrationVideoViewTest.java
+++ b/library/test/robotest/src/com/android/setupwizardlib/view/IllustrationVideoViewTest.java
@@ -48,7 +48,7 @@ import org.robolectric.annotation.Config;
 import org.robolectric.annotation.Implementation;
 import org.robolectric.annotation.Implements;
 import org.robolectric.annotation.RealObject;
-import org.robolectric.internal.Shadow;
+import org.robolectric.shadow.api.Shadow;
 import org.robolectric.shadows.ShadowMediaPlayer;
 import org.robolectric.util.ReflectionHelpers;