2 files changed, 54 insertions, 3 deletions

diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowManagerServiceRule.java b/services/tests/servicestests/src/com/android/server/wm/WindowManagerServiceRule.java
index 54fd7db9b537..cff9504ef9dc 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowManagerServiceRule.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowManagerServiceRule.java
@@ -34,8 +34,8 @@ import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManagerInternal;
 import android.os.PowerSaveState;
 import android.view.InputChannel;
-
-import androidx.test.InstrumentationRegistry;
+import android.view.SurfaceControl;
+import android.view.SurfaceControl.Transaction;
 
 import com.android.server.LocalServices;
 import com.android.server.input.InputManagerService;
@@ -46,6 +46,12 @@ import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 import org.mockito.invocation.InvocationOnMock;
 
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.List;
+
+import androidx.test.InstrumentationRegistry;
+
 /**
  * A test rule that sets up a fresh WindowManagerService instance before each test and makes sure
  * to properly tear it down after.
@@ -61,6 +67,12 @@ public class WindowManagerServiceRule implements TestRule {
 
     private WindowManagerService mService;
     private TestWindowManagerPolicy mPolicy;
+    // Record all {@link SurfaceControl.Transaction} created while testing and releases native
+    // resources when test finishes.
+    private final List<WeakReference<Transaction>> mSurfaceTransactions = new ArrayList<>();
+    // Record all {@link SurfaceControl} created while testing and releases native resources when
+    // test finishes.
+    private final List<WeakReference<SurfaceControl>> mSurfaceControls = new ArrayList<>();
 
     @Override
     public Statement apply(Statement base, Description description) {
@@ -114,6 +126,19 @@ public class WindowManagerServiceRule implements TestRule {
                 mService = WindowManagerService.main(context, ims, false,
                         false, mPolicy = new TestWindowManagerPolicy(
                                 WindowManagerServiceRule.this::getWindowManagerService));
+                mService.mTransactionFactory = () -> {
+                    final SurfaceControl.Transaction transaction = new SurfaceControl.Transaction();
+                    mSurfaceTransactions.add(new WeakReference<>(transaction));
+                    return transaction;
+                };
+                mService.mSurfaceBuilderFactory = session -> new SurfaceControl.Builder(session) {
+                    @Override
+                    public SurfaceControl build() {
+                        final SurfaceControl control = super.build();
+                        mSurfaceControls.add(new WeakReference<>(control));
+                        return control;
+                    }
+                };
 
                 mService.onInitReady();
 
@@ -135,6 +160,8 @@ public class WindowManagerServiceRule implements TestRule {
 
             private void tearDown() {
                 waitUntilWindowManagerHandlersIdle();
+                destroyAllSurfaceTransactions();
+                destroyAllSurfaceControls();
                 removeServices();
                 mService = null;
                 mPolicy = null;
@@ -158,4 +185,24 @@ public class WindowManagerServiceRule implements TestRule {
             SurfaceAnimationThread.getHandler().runWithScissors(() -> { }, 0);
         }
     }
+
+    private void destroyAllSurfaceTransactions() {
+        for (final WeakReference<Transaction> reference : mSurfaceTransactions) {
+            final Transaction transaction = reference.get();
+            if (transaction != null) {
+                reference.clear();
+                transaction.close();
+            }
+        }
+    }
+
+    private void destroyAllSurfaceControls() {
+        for (final WeakReference<SurfaceControl> reference : mSurfaceControls) {
+            final SurfaceControl control = reference.get();
+            if (control != null) {
+                reference.clear();
+                control.destroy();
+            }
+        }
+    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java b/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java
index a610e6efda8c..3a8c4ae73493 100644
--- a/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java
@@ -98,10 +98,13 @@ public class ZOrderingTests extends WindowTestsBase {
             super(s);
         }
 
+        @Override
         public SurfaceControl.Builder setParent(SurfaceControl sc) {
             mPendingParent = sc;
             return super.setParent(sc);
         }
+
+        @Override
         public SurfaceControl build() {
             SurfaceControl sc = super.build();
             mParentFor.put(sc, mPendingParent);
@@ -110,7 +113,7 @@ public class ZOrderingTests extends WindowTestsBase {
         }
     }
 
-    class HierarchyRecordingBuilderFactory implements SurfaceBuilderFactory {
+    private class HierarchyRecordingBuilderFactory implements SurfaceBuilderFactory {
         public SurfaceControl.Builder make(SurfaceSession s) {
             return new HierarchyRecorder(s);
         }
@@ -131,6 +134,7 @@ public class ZOrderingTests extends WindowTestsBase {
     @After
     public void after() {
         mTransaction.close();
+        mParentFor.keySet().forEach(SurfaceControl::destroy);
         mParentFor.clear();
     }