diff --git a/core/java/android/view/WindowManagerImpl.java b/core/java/android/view/WindowManagerImpl.java
index 316a5f2c88d21..561ee604aa7fe 100644
--- a/core/java/android/view/WindowManagerImpl.java
+++ b/core/java/android/view/WindowManagerImpl.java
@@ -36,6 +36,7 @@ import android.os.Bundle;
 import android.os.IBinder;
 import android.os.RemoteException;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.IResultReceiver;
 
 import java.util.List;
@@ -69,7 +70,8 @@ import java.util.List;
 public final class WindowManagerImpl implements WindowManager {
     @UnsupportedAppUsage
     private final WindowManagerGlobal mGlobal = WindowManagerGlobal.getInstance();
-    private final Context mContext;
+    @VisibleForTesting
+    public final Context mContext;
     private final Window mParentWindow;
 
     private IBinder mDefaultToken;
diff --git a/core/java/com/android/internal/policy/DecorContext.java b/core/java/com/android/internal/policy/DecorContext.java
index 99b4b5fb77077..f00776897f2c8 100644
--- a/core/java/com/android/internal/policy/DecorContext.java
+++ b/core/java/com/android/internal/policy/DecorContext.java
@@ -41,17 +41,17 @@ import java.lang.ref.WeakReference;
 public class DecorContext extends ContextThemeWrapper {
     private PhoneWindow mPhoneWindow;
     private WindowManager mWindowManager;
-    private Resources mActivityResources;
+    private Resources mResources;
     private ContentCaptureManager mContentCaptureManager;
 
-    private WeakReference<Context> mActivityContext;
+    private WeakReference<Context> mContext;
 
     // TODO(b/149928768): Non-activity context can be passed.
     @VisibleForTesting
-    public DecorContext(Context context, Context activityContext) {
-        super(context.createDisplayContext(activityContext.getDisplayNoVerify()), null);
-        mActivityContext = new WeakReference<>(activityContext);
-        mActivityResources = activityContext.getResources();
+    public DecorContext(Context baseContext, Context context) {
+        super(baseContext.createDisplayContext(context.getDisplayNoVerify()), null);
+        mContext = new WeakReference<>(context);
+        mResources = context.getResources();
     }
 
     void setPhoneWindow(PhoneWindow phoneWindow) {
@@ -61,58 +61,56 @@ public class DecorContext extends ContextThemeWrapper {
 
     @Override
     public Object getSystemService(String name) {
+        final Context context = mContext.get();
         if (Context.WINDOW_SERVICE.equals(name)) {
-            if (mWindowManager == null) {
-                WindowManagerImpl wm =
-                        (WindowManagerImpl) super.getSystemService(Context.WINDOW_SERVICE);
+            if (context != null && mWindowManager == null) {
+                WindowManagerImpl wm = (WindowManagerImpl) context.getSystemService(name);
                 mWindowManager = wm.createLocalWindowManager(mPhoneWindow);
             }
             return mWindowManager;
         }
         if (Context.CONTENT_CAPTURE_MANAGER_SERVICE.equals(name)) {
-            if (mContentCaptureManager == null) {
-                Context activityContext = mActivityContext.get();
-                if (activityContext != null) {
-                    mContentCaptureManager = (ContentCaptureManager) activityContext
-                            .getSystemService(name);
-                }
+            if (context != null && mContentCaptureManager == null) {
+                mContentCaptureManager = (ContentCaptureManager) context.getSystemService(name);
             }
             return mContentCaptureManager;
         }
-        return super.getSystemService(name);
+        // LayoutInflater and WallpaperManagerService should also be obtained from context
+        // instead of application context.
+        return (context != null) ? context.getSystemService(name) : super.getSystemService(name);
     }
 
     @Override
     public Resources getResources() {
-        Context activityContext = mActivityContext.get();
+        Context context = mContext.get();
         // Attempt to update the local cached Resources from the activity context. If the activity
         // is no longer around, return the old cached values.
-        if (activityContext != null) {
-            mActivityResources = activityContext.getResources();
+        if (context != null) {
+            mResources = context.getResources();
         }
 
-        return mActivityResources;
+        return mResources;
     }
 
     @Override
     public AssetManager getAssets() {
-        return mActivityResources.getAssets();
+        return mResources.getAssets();
     }
 
     @Override
     public AutofillOptions getAutofillOptions() {
-        Context activityContext = mActivityContext.get();
-        if (activityContext != null) {
-            return activityContext.getAutofillOptions();
+        Context context = mContext.get();
+        if (context != null) {
+            return context.getAutofillOptions();
         }
         return null;
     }
 
     @Override
     public ContentCaptureOptions getContentCaptureOptions() {
-        Context activityContext = mActivityContext.get();
-        if (activityContext != null) {
-            return activityContext.getContentCaptureOptions();
+        Context context = mContext.get();
+        if (context != null) {
+            return context.getContentCaptureOptions();
         }
         return null;
     }
diff --git a/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java b/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
index 3e40466e4b649..d019704fb6844 100644
--- a/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
+++ b/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
@@ -20,19 +20,24 @@ import static android.view.Display.DEFAULT_DISPLAY;
 
 import static org.junit.Assert.assertEquals;
 
+import android.app.Activity;
+import android.app.EmptyActivity;
 import android.content.Context;
 import android.hardware.display.DisplayManagerGlobal;
 import android.platform.test.annotations.Presubmit;
 import android.view.Display;
 import android.view.DisplayAdjustments;
 import android.view.DisplayInfo;
+import android.view.WindowManager;
+import android.view.WindowManagerImpl;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.core.app.ApplicationProvider;
 import androidx.test.filters.SmallTest;
+import androidx.test.rule.ActivityTestRule;
 import androidx.test.runner.AndroidJUnit4;
 
-
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -46,9 +51,13 @@ public final class DecorContextTest {
     private Context mContext;
     private static final int EXTERNAL_DISPLAY = DEFAULT_DISPLAY + 1;
 
+    @Rule
+    public ActivityTestRule<EmptyActivity> mActivityRule =
+            new ActivityTestRule<>(EmptyActivity.class);
+
     @Before
-    public void setUp() throws Exception {
-        mContext = InstrumentationRegistry.getContext();
+    public void setUp() {
+        mContext = ApplicationProvider.getApplicationContext();
     }
 
     @Test
@@ -76,4 +85,19 @@ public final class DecorContextTest {
         Display associatedDisplay = decorContext.getDisplay();
         assertEquals(expectedDisplayId, associatedDisplay.getDisplayId());
     }
+
+    @Test
+    public void testGetWindowManagerFromVisualDecorContext() throws Throwable {
+        mActivityRule.runOnUiThread(() -> {
+            Activity activity = mActivityRule.getActivity();
+            final DecorContext decorContext = new DecorContext(mContext.getApplicationContext(),
+                    activity);
+            WindowManagerImpl actualWm = (WindowManagerImpl)
+                    decorContext.getSystemService(WindowManager.class);
+            WindowManagerImpl expectedWm = (WindowManagerImpl)
+                    activity.getSystemService(WindowManager.class);
+            // Verify that window manager is from activity not application context.
+            assertEquals(expectedWm.mContext, actualWm.mContext);
+        });
+    }
 }