diff --git a/tests/unit/src/com/android/settings/communal/CommunalPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/communal/CommunalPreferenceControllerTest.java
similarity index 53%
rename from tests/unit/src/com/android/settings/communal/CommunalPreferenceControllerTest.java
rename to tests/robotests/src/com/android/settings/communal/CommunalPreferenceControllerTest.java
index 6bd84f0064b..2114123e62f 100644
--- a/tests/unit/src/com/android/settings/communal/CommunalPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/communal/CommunalPreferenceControllerTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * Copyright (C) 2024 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,75 +18,63 @@ package com.android.settings.communal;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.when;
 
 import android.content.Context;
-import android.content.res.Resources;
-import android.os.UserManager;
+import android.os.UserHandle;
 
-import androidx.test.core.app.ApplicationProvider;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-
-import com.android.settings.Utils;
-import com.android.settings.testutils.ResourcesUtils;
+import com.android.settings.R;
+import com.android.settings.testutils.shadow.SettingsShadowResources;
+import com.android.settings.testutils.shadow.ShadowUserManager;
 
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
 
-@RunWith(AndroidJUnit4.class)
-@Ignore("b/337417918")
+@RunWith(RobolectricTestRunner.class)
+@Config(shadows = {SettingsShadowResources.class, ShadowUserManager.class})
 public class CommunalPreferenceControllerTest {
-    @Mock
-    private UserManager mUserManager;
-
-    private Context mContext;
-    private Resources mResources;
+    private ShadowUserManager mShadowUserManager;
     private CommunalPreferenceController mController;
 
     private static final String PREF_KEY = "test_key";
 
     @Before
     public void setup() {
-        MockitoAnnotations.initMocks(this);
-        mContext = spy(ApplicationProvider.getApplicationContext());
-        mResources = spy(mContext.getResources());
-
-        mController = new CommunalPreferenceController(mContext, PREF_KEY);
-
-        when(mContext.getResources()).thenReturn(mResources);
-        when(mContext.getSystemService(UserManager.class)).thenReturn(mUserManager);
+        final Context context = spy(RuntimeEnvironment.application);
+        mShadowUserManager = ShadowUserManager.getShadow();
+        doReturn(context).when(context).createContextAsUser(any(UserHandle.class), anyInt());
+        mController = new CommunalPreferenceController(context, PREF_KEY);
     }
 
     @Test
     public void isAvailable_communalEnabled_shouldBeTrueForDockUser() {
         setCommunalEnabled(true);
-        when(Utils.canCurrentUserDream(mContext)).thenReturn(true);
+        mShadowUserManager.setUserForeground(true);
         assertTrue(mController.isAvailable());
     }
 
     @Test
     public void isAvailable_communalEnabled_shouldBeFalseForNonDockUser() {
         setCommunalEnabled(true);
-        when(Utils.canCurrentUserDream(mContext)).thenReturn(false);
+        mShadowUserManager.setUserForeground(false);
         assertFalse(mController.isAvailable());
     }
 
     @Test
     public void isAvailable_communalDisabled_shouldBeFalseForDockUser() {
         setCommunalEnabled(false);
-        when(Utils.canCurrentUserDream(mContext)).thenReturn(true);
+        mShadowUserManager.setUserForeground(true);
         assertFalse(mController.isAvailable());
     }
 
     private void setCommunalEnabled(boolean enabled) {
-        final int boolId = ResourcesUtils.getResourcesId(
-                ApplicationProvider.getApplicationContext(), "bool",
-                "config_show_communal_settings");
-        when(mResources.getBoolean(boolId)).thenReturn(enabled);
+        SettingsShadowResources.overrideResource(R.bool.config_show_communal_settings, enabled);
     }
 }
diff --git a/tests/robotests/testutils/com/android/settings/testutils/shadow/ShadowUserManager.java b/tests/robotests/testutils/com/android/settings/testutils/shadow/ShadowUserManager.java
index dfcb7ee71aa..417b102dcc9 100644
--- a/tests/robotests/testutils/com/android/settings/testutils/shadow/ShadowUserManager.java
+++ b/tests/robotests/testutils/com/android/settings/testutils/shadow/ShadowUserManager.java
@@ -309,6 +309,11 @@ public class ShadowUserManager extends org.robolectric.shadows.ShadowUserManager
                 UserInfo.FLAG_INITIALIZED | UserInfo.FLAG_ADMIN | UserInfo.FLAG_PRIMARY);
     }
 
+    @Implementation
+    protected UserHandle getMainUser() {
+        return UserHandle.of(PRIMARY_USER_ID);
+    }
+
     protected boolean setUserEphemeral(@UserIdInt int userId, boolean enableEphemeral) {
         UserInfo userInfo = mUserProfileInfos.stream()
                 .filter(user -> user.id == userId)