diff --git a/services/core/java/com/android/server/wm/ScreenRotationAnimation.java b/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
index 84cd8d1..2d5c97f 100644
--- a/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
+++ b/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
@@ -257,7 +257,7 @@
         mOriginalWidth = originalWidth;
         mOriginalHeight = originalHeight;
 
-        final SurfaceControl.Transaction t = new SurfaceControl.Transaction();
+        final SurfaceControl.Transaction t = mService.mTransactionFactory.make();
         try {
             mSurfaceControl = displayContent.makeOverlay()
                     .setName("ScreenshotSurface")
@@ -267,13 +267,13 @@
 
             // In case display bounds change, screenshot buffer and surface may mismatch so set a
             // scaling mode.
-            SurfaceControl.Transaction t2 = new SurfaceControl.Transaction();
+            SurfaceControl.Transaction t2 = mService.mTransactionFactory.make();
             t2.setOverrideScalingMode(mSurfaceControl, Surface.SCALING_MODE_SCALE_TO_WINDOW);
             t2.apply(true /* sync */);
 
             // Capture a screenshot into the surface we just created.
             final int displayId = display.getDisplayId();
-            final Surface surface = new Surface();
+            final Surface surface = mService.mSurfaceFactory.make();
             surface.copyFrom(mSurfaceControl);
             if (mService.mDisplayManagerInternal.screenshot(displayId, surface)) {
                 t.setLayer(mSurfaceControl, SCREEN_FREEZE_LAYER_SCREENSHOT);
diff --git a/services/core/java/com/android/server/wm/SurfaceFactory.java b/services/core/java/com/android/server/wm/SurfaceFactory.java
new file mode 100644
index 0000000..076b7df
--- /dev/null
+++ b/services/core/java/com/android/server/wm/SurfaceFactory.java
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wm;
+
+import android.view.Surface;
+
+/**
+ * Helper class to inject custom {@link Surface} objects into window manager.
+ */
+interface SurfaceFactory {
+    Surface make();
+};
+
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index 74fb3fa..dddc6b7 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -50,6 +50,7 @@
 import android.view.WindowManager;
 import android.view.animation.Animation;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ToBooleanFunction;
 
 import java.io.PrintWriter;
@@ -700,24 +701,37 @@
         mWallpaperTokens.remove(token);
     }
 
+
+    @VisibleForTesting
+    boolean canScreenshotWallpaper() {
+        return canScreenshotWallpaper(getTopVisibleWallpaper());
+    }
+
+    private boolean canScreenshotWallpaper(WindowState wallpaperWindowState) {
+        if (!mService.mPolicy.isScreenOn()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
+            }
+            return false;
+        }
+
+        if (wallpaperWindowState == null) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "No visible wallpaper to screenshot");
+            }
+            return false;
+        }
+        return true;
+    }
+
     /**
      * Take a screenshot of the wallpaper if it's visible.
      *
      * @return Bitmap of the wallpaper
      */
     Bitmap screenshotWallpaperLocked() {
-        if (!mService.mPolicy.isScreenOn()) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
-            }
-            return null;
-        }
-
         final WindowState wallpaperWindowState = getTopVisibleWallpaper();
-        if (wallpaperWindowState == null) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.i(TAG_WM, "No visible wallpaper to screenshot");
-            }
+        if (!canScreenshotWallpaper(wallpaperWindowState)) {
             return null;
         }
 
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 39df0e4..d4430da4 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -814,8 +814,9 @@
 
     SurfaceBuilderFactory mSurfaceBuilderFactory = SurfaceControl.Builder::new;
     TransactionFactory mTransactionFactory = SurfaceControl.Transaction::new;
+    SurfaceFactory mSurfaceFactory = Surface::new;
 
-    private final SurfaceControl.Transaction mTransaction = mTransactionFactory.make();
+    private final SurfaceControl.Transaction mTransaction;
 
     static void boostPriorityForLockedSection() {
         sThreadPriorityBooster.boost();
@@ -909,9 +910,21 @@
     public static WindowManagerService main(final Context context, final InputManagerService im,
             final boolean showBootMsgs, final boolean onlyCore, WindowManagerPolicy policy,
             ActivityTaskManagerService atm) {
+        return main(context, im, showBootMsgs, onlyCore, policy, atm,
+                SurfaceControl.Transaction::new);
+    }
+
+    /**
+     * Creates and returns an instance of the WindowManagerService. This call allows the caller
+     * to override the {@link TransactionFactory} to stub functionality under test.
+     */
+    @VisibleForTesting
+    public static WindowManagerService main(final Context context, final InputManagerService im,
+            final boolean showBootMsgs, final boolean onlyCore, WindowManagerPolicy policy,
+            ActivityTaskManagerService atm, TransactionFactory transactionFactory) {
         DisplayThread.getHandler().runWithScissors(() ->
                 sInstance = new WindowManagerService(context, im, showBootMsgs, onlyCore, policy,
-                        atm), 0);
+                        atm, transactionFactory), 0);
         return sInstance;
     }
 
@@ -933,7 +946,7 @@
 
     private WindowManagerService(Context context, InputManagerService inputManager,
             boolean showBootMsgs, boolean onlyCore, WindowManagerPolicy policy,
-            ActivityTaskManagerService atm) {
+            ActivityTaskManagerService atm, TransactionFactory transactionFactory) {
         installLock(this, INDEX_WINDOW);
         mGlobalLock = atm.getGlobalLock();
         mAtmService = atm;
@@ -962,6 +975,9 @@
         mDisplayManagerInternal = LocalServices.getService(DisplayManagerInternal.class);
         mDisplayWindowSettings = new DisplayWindowSettings(this);
 
+
+        mTransactionFactory = transactionFactory;
+        mTransaction = mTransactionFactory.make();
         mPolicy = policy;
         mAnimator = new WindowAnimator(this);
         mRoot = new RootWindowContainer(this);
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppChangeTransitionTests.java b/services/tests/wmtests/src/com/android/server/wm/AppChangeTransitionTests.java
index 2627ec7..c072d4e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppChangeTransitionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppChangeTransitionTests.java
@@ -59,8 +59,7 @@
     public void setUpOnDisplay(DisplayContent dc) {
         mStack = createTaskStackOnDisplay(WINDOWING_MODE_UNDEFINED, ACTIVITY_TYPE_STANDARD, dc);
         mTask = createTaskInStack(mStack, 0 /* userId */);
-        mToken = WindowTestUtils.createTestAppWindowToken(dc);
-        mToken.mSkipOnParentChanged = false;
+        mToken = WindowTestUtils.createTestAppWindowToken(dc, false /* skipOnParentChanged */);
 
         mTask.addChild(mToken, 0);
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenAnimationTests.java b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenAnimationTests.java
index 108ee180b..e007c86 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenAnimationTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenAnimationTests.java
@@ -18,7 +18,6 @@
 
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
-import static android.view.SurfaceControl.Transaction;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 
@@ -30,6 +29,9 @@
 import android.platform.test.annotations.Presubmit;
 import android.view.SurfaceControl;
 
+import androidx.test.filters.FlakyTest;
+import androidx.test.filters.SmallTest;
+
 import com.android.server.wm.WindowTestUtils.TestAppWindowToken;
 
 import org.junit.Before;
@@ -38,8 +40,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import androidx.test.filters.FlakyTest;
-import androidx.test.filters.SmallTest;
 
 /**
  * Animation related tests for the {@link AppWindowToken} class.
@@ -55,8 +55,6 @@
     private TestAppWindowToken mToken;
 
     @Mock
-    private Transaction mTransaction;
-    @Mock
     private AnimationAdapter mSpec;
 
     @Before
@@ -65,7 +63,6 @@
 
         mToken = createTestAppWindowToken(mDisplayContent, WINDOWING_MODE_FULLSCREEN,
                 ACTIVITY_TYPE_STANDARD, false /* skipOnParentChanged */);
-        mToken.setPendingTransaction(mTransaction);
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
index 1dd72ec..2c575f5 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
@@ -60,7 +60,7 @@
  * Tests for the {@link AppWindowToken} class.
  *
  * Build/Install/Run:
- *  atest FrameworksServicesTests:AppWindowTokenTests
+ *  atest WmTests:AppWindowTokenTests
  */
 @SmallTest
 @Presubmit
@@ -76,9 +76,9 @@
     public void setUp() throws Exception {
         mStack = createTaskStackOnDisplay(mDisplayContent);
         mTask = createTaskInStack(mStack, 0 /* userId */);
-        mToken = WindowTestUtils.createTestAppWindowToken(mDisplayContent);
+        mToken = WindowTestUtils.createTestAppWindowToken(mDisplayContent,
+                false /* skipOnParentChanged */);
 
-        mToken.mSkipOnParentChanged = false;
         mTask.addChild(mToken, 0);
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/MockSurfaceControlBuilder.java b/services/tests/wmtests/src/com/android/server/wm/MockSurfaceControlBuilder.java
new file mode 100644
index 0000000..66139e6
--- /dev/null
+++ b/services/tests/wmtests/src/com/android/server/wm/MockSurfaceControlBuilder.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wm;
+
+import static org.mockito.Mockito.mock;
+
+import android.view.SurfaceControl;
+
+/**
+ * Stubbed {@link SurfaceControl.Builder} class that returns a mocked SurfaceControl instance
+ * that can be used for unit testing.
+ */
+class MockSurfaceControlBuilder extends SurfaceControl.Builder {
+    @Override
+    public SurfaceControl.Builder setParent(SurfaceControl sc) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl build() {
+        return mock(SurfaceControl.class);
+    }
+}
diff --git a/services/tests/wmtests/src/com/android/server/wm/StubTransaction.java b/services/tests/wmtests/src/com/android/server/wm/StubTransaction.java
new file mode 100644
index 0000000..d919fc8
--- /dev/null
+++ b/services/tests/wmtests/src/com/android/server/wm/StubTransaction.java
@@ -0,0 +1,251 @@
+/*
+ * Copyright (C) 2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wm;
+
+import android.graphics.Matrix;
+import android.graphics.Rect;
+import android.graphics.Region;
+import android.os.IBinder;
+import android.os.Parcel;
+import android.view.InputWindowHandle;
+import android.view.Surface;
+import android.view.SurfaceControl;
+
+/**
+ * Stubbed {@link android.view.SurfaceControl.Transaction} class that can be used when unit
+ * testing to avoid calls to native code.
+ */
+public class StubTransaction extends SurfaceControl.Transaction {
+    @Override
+    public void apply() {
+    }
+
+    @Override
+    public void close() {
+    }
+
+    @Override
+    public void apply(boolean sync) {
+    }
+
+    @Override
+    public SurfaceControl.Transaction setVisibility(SurfaceControl sc, boolean visible) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction show(SurfaceControl sc) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction hide(SurfaceControl sc) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setPosition(SurfaceControl sc, float x, float y) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setBufferSize(SurfaceControl sc,
+            int w, int h) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setLayer(SurfaceControl sc, int z) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setRelativeLayer(SurfaceControl sc, SurfaceControl relativeTo,
+            int z) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setTransparentRegionHint(SurfaceControl sc,
+            Region transparentRegion) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setAlpha(SurfaceControl sc, float alpha) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setInputWindowInfo(SurfaceControl sc,
+            InputWindowHandle handle) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction transferTouchFocus(IBinder fromToken, IBinder toToken) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setGeometry(SurfaceControl sc, Rect sourceCrop,
+            Rect destFrame, @Surface.Rotation int orientation) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setMatrix(SurfaceControl sc,
+            float dsdx, float dtdx, float dtdy, float dsdy) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setMatrix(SurfaceControl sc, Matrix matrix, float[] float9) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setColorTransform(SurfaceControl sc, float[] matrix,
+            float[] translation) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setWindowCrop(SurfaceControl sc, Rect crop) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setWindowCrop(SurfaceControl sc, int width, int height) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setCornerRadius(SurfaceControl sc, float cornerRadius) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setLayerStack(SurfaceControl sc, int layerStack) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction deferTransactionUntil(SurfaceControl sc, IBinder handle,
+            long frameNumber) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction deferTransactionUntilSurface(SurfaceControl sc,
+            Surface barrierSurface,
+            long frameNumber) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction reparentChildren(SurfaceControl sc, IBinder newParentHandle) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction reparent(SurfaceControl sc, SurfaceControl newParent) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction detachChildren(SurfaceControl sc) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setOverrideScalingMode(SurfaceControl sc,
+            int overrideScalingMode) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setColor(SurfaceControl sc, float[] color) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setGeometryAppliesWithResize(SurfaceControl sc) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setSecure(SurfaceControl sc, boolean isSecure) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setOpaque(SurfaceControl sc, boolean isOpaque) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setDisplaySurface(IBinder displayToken, Surface surface) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setDisplayLayerStack(IBinder displayToken, int layerStack) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setDisplayProjection(IBinder displayToken,
+            int orientation, Rect layerStackRect, Rect displayRect) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setDisplaySize(IBinder displayToken, int width, int height) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setAnimationTransaction() {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setEarlyWakeup() {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setMetadata(int key, int data) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction setMetadata(int key, Parcel data) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction merge(SurfaceControl.Transaction other) {
+        return this;
+    }
+
+    @Override
+    public SurfaceControl.Transaction remove(SurfaceControl sc) {
+        return this;
+    }
+}
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index 712cd1e..366acea 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -167,7 +167,7 @@
 
         mWindowManagerPolicy = new TestWindowManagerPolicy(this::getWindowManagerService);
         mWindowManagerService = WindowManagerService.main(
-                context, ims, false, false, mWindowManagerPolicy, atms);
+                context, ims, false, false, mWindowManagerPolicy, atms, StubTransaction::new);
 
         mWindowManagerService.onInitReady();
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
index d07230e..6249bde 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
@@ -21,11 +21,9 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
 
-import static junit.framework.TestCase.assertNotNull;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
-import static org.junit.Assert.assertNull;
-
-import android.graphics.Bitmap;
 import android.os.IBinder;
 import android.platform.test.annotations.Presubmit;
 
@@ -37,7 +35,7 @@
  * Tests for the {@link WallpaperController} class.
  *
  * Build/Install/Run:
- *  atest FrameworksServicesTests:WallpaperControllerTests
+ *  atest WmTests:WallpaperControllerTests
  */
 @SmallTest
 @Presubmit
@@ -49,34 +47,29 @@
         synchronized (mWm.mGlobalLock) {
             // No wallpaper
             final DisplayContent dc = createNewDisplay();
-            Bitmap wallpaperBitmap = dc.mWallpaperController.screenshotWallpaperLocked();
-            assertNull(wallpaperBitmap);
+            assertFalse(dc.mWallpaperController.canScreenshotWallpaper());
 
             // No wallpaper WSA Surface
             WindowToken wallpaperWindowToken = new WallpaperWindowToken(mWm, mock(IBinder.class),
                     true, dc, true /* ownerCanManageAppTokens */);
             WindowState wallpaperWindow = createWindow(null /* parent */, TYPE_WALLPAPER,
                     wallpaperWindowToken, "wallpaperWindow");
-            wallpaperBitmap = dc.mWallpaperController.screenshotWallpaperLocked();
-            assertNull(wallpaperBitmap);
+            assertFalse(dc.mWallpaperController.canScreenshotWallpaper());
 
             // Wallpaper with not visible WSA surface.
             wallpaperWindow.mWinAnimator.mSurfaceController = windowSurfaceController;
             wallpaperWindow.mWinAnimator.mLastAlpha = 1;
-            wallpaperBitmap = dc.mWallpaperController.screenshotWallpaperLocked();
-            assertNull(wallpaperBitmap);
+            assertFalse(dc.mWallpaperController.canScreenshotWallpaper());
 
             when(windowSurfaceController.getShown()).thenReturn(true);
 
             // Wallpaper with WSA alpha set to 0.
             wallpaperWindow.mWinAnimator.mLastAlpha = 0;
-            wallpaperBitmap = dc.mWallpaperController.screenshotWallpaperLocked();
-            assertNull(wallpaperBitmap);
+            assertFalse(dc.mWallpaperController.canScreenshotWallpaper());
 
             // Wallpaper window with WSA Surface
             wallpaperWindow.mWinAnimator.mLastAlpha = 1;
-            wallpaperBitmap = dc.mWallpaperController.screenshotWallpaperLocked();
-            assertNotNull(wallpaperBitmap);
+            assertTrue(dc.mWallpaperController.canScreenshotWallpaper());
         }
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java
index 5bfa0c6..da1defa 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestUtils.java
@@ -28,7 +28,6 @@
 import android.os.IBinder;
 import android.view.IApplicationToken;
 import android.view.IWindow;
-import android.view.SurfaceControl.Transaction;
 import android.view.WindowManager;
 
 /**
@@ -72,7 +71,6 @@
     static class TestAppWindowToken extends AppWindowToken {
         boolean mOnTop = false;
         private boolean mSkipPrepareSurfaces;
-        private Transaction mPendingTransactionOverride;
         boolean mSkipOnParentChanged = true;
 
         private TestAppWindowToken(DisplayContent dc, boolean skipOnParentChanged) {
@@ -126,17 +124,6 @@
         void setSkipPrepareSurfaces(boolean ignore) {
             mSkipPrepareSurfaces = ignore;
         }
-
-        void setPendingTransaction(Transaction transaction) {
-            mPendingTransactionOverride = transaction;
-        }
-
-        @Override
-        public Transaction getPendingTransaction() {
-            return mPendingTransactionOverride == null
-                    ? super.getPendingTransaction()
-                    : mPendingTransactionOverride;
-        }
     }
 
     /**
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index 0d7d085..d202e16 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -37,7 +37,10 @@
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
+
+import static org.mockito.Mockito.mock;
+
 
 import android.content.Context;
 import android.content.res.Configuration;
@@ -47,6 +50,8 @@
 import android.view.Display;
 import android.view.DisplayInfo;
 import android.view.IWindow;
+import android.view.Surface;
+import android.view.SurfaceControl.Transaction;
 import android.view.WindowManager;
 
 import com.android.server.AttributeCache;
@@ -96,10 +101,9 @@
     private MockTracker mMockTracker;
 
     /**
-     * To restore the original SurfaceControl.Transaction factory if any tests changed
-     * {@link WindowManagerService#mTransactionFactory}.
+     * Spied {@link Transaction} class than can be used to verify calls.
      */
-    private TransactionFactory mOriginalTransactionFactory;
+    Transaction mTransaction;
 
     @Rule
     public final DexmakerShareClassLoaderRule mDexmakerShareClassLoaderRule =
@@ -129,11 +133,21 @@
         // in the set up are clear. This can be removed when b/37850063 is fixed.
         try {
             mMockSession = mock(Session.class);
+            mTransaction = spy(StubTransaction.class);
 
             final Context context = getInstrumentation().getTargetContext();
 
             mWm = mSystemServicesTestRule.getWindowManagerService();
-            mOriginalTransactionFactory = mWm.mTransactionFactory;
+
+            // Setup factory classes to prevent calls to native code.
+
+            // Return a spied Transaction class than can be used to verify calls.
+            mWm.mTransactionFactory = () -> mTransaction;
+            // Return a SurfaceControl.Builder class that creates mocked SurfaceControl instances.
+            mWm.mSurfaceBuilderFactory = (unused) -> new MockSurfaceControlBuilder();
+            // Return mocked Surface instances.
+            mWm.mSurfaceFactory = () -> mock(Surface.class);
+
             beforeCreateDisplay();
 
             context.getDisplay().getDisplayInfo(mDisplayInfo);
@@ -183,7 +197,6 @@
             // stable state to clean up for consistency.
             waitUntilHandlersIdle();
 
-            mWm.mTransactionFactory = mOriginalTransactionFactory;
             final LinkedList<WindowState> nonCommonWindows = new LinkedList<>();
 
             synchronized (mWm.mGlobalLock) {