diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index a4256a9d8ab3e..48160e475d504 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -2747,6 +2747,7 @@ class ContextImpl extends Context {
             opPackageName = container.mOpPackageName;
             setResources(container.mResources);
             mDisplay = container.mDisplay;
+            mIsAssociatedWithDisplay = container.mIsAssociatedWithDisplay;
             mIsSystemOrSystemUiContext = container.mIsSystemOrSystemUiContext;
         } else {
             mBasePackageName = packageInfo.mPackageName;
diff --git a/core/tests/coretests/src/android/content/ContextTest.java b/core/tests/coretests/src/android/content/ContextTest.java
index 2057a8181147e..f13a11e9edfba 100644
--- a/core/tests/coretests/src/android/content/ContextTest.java
+++ b/core/tests/coretests/src/android/content/ContextTest.java
@@ -16,6 +16,8 @@
 
 package android.content;
 
+import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY;
+import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -23,8 +25,13 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 
 import android.app.ActivityThread;
+import android.content.res.Configuration;
+import android.graphics.PixelFormat;
 import android.hardware.display.DisplayManager;
+import android.hardware.display.VirtualDisplay;
+import android.media.ImageReader;
 import android.os.UserHandle;
+import android.view.Display;
 
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -127,4 +134,41 @@ public class ContextTest {
 
         assertThat(systemUiContext.isUiContext()).isTrue();
     }
+
+    @Test
+    public void testGetDisplayFromDisplayContextDerivedContextOnPrimaryDisplay() {
+        verifyGetDisplayFromDisplayContextDerivedContext(false /* onSecondaryDisplay */);
+    }
+
+    @Test
+    public void testGetDisplayFromDisplayContextDerivedContextOnSecondaryDisplay() {
+        verifyGetDisplayFromDisplayContextDerivedContext(true /* onSecondaryDisplay */);
+    }
+
+    private static void verifyGetDisplayFromDisplayContextDerivedContext(
+            boolean onSecondaryDisplay) {
+        final Context appContext = ApplicationProvider.getApplicationContext();
+        final DisplayManager displayManager = appContext.getSystemService(DisplayManager.class);
+        final Display display;
+        if (onSecondaryDisplay) {
+            display = getSecondaryDisplay(displayManager);
+        } else {
+            display = displayManager.getDisplay(DEFAULT_DISPLAY);
+        }
+        final Context context = appContext.createDisplayContext(display)
+                .createConfigurationContext(new Configuration());
+        assertEquals(display, context.getDisplay());
+    }
+
+    private static Display getSecondaryDisplay(DisplayManager displayManager) {
+        final int width = 800;
+        final int height = 480;
+        final int density = 160;
+        ImageReader reader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888,
+                2 /* maxImages */);
+        VirtualDisplay virtualDisplay = displayManager.createVirtualDisplay(
+                ContextTest.class.getName(), width, height, density, reader.getSurface(),
+                VIRTUAL_DISPLAY_FLAG_PUBLIC | VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY);
+        return virtualDisplay.getDisplay();
+    }
 }