diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java b/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
index 4104999..12b7445 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
@@ -16,6 +16,8 @@
 
 package com.android.server.wm;
 
+import static android.tools.traces.Utils.busyWaitForDataSourceRegistration;
+
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verifyZeroInteractions;
@@ -28,7 +30,6 @@
 import static java.io.File.createTempFile;
 import static java.nio.file.Files.createTempDirectory;
 
-import android.os.ParcelFileDescriptor;
 import android.platform.test.annotations.Presubmit;
 import android.tools.ScenarioBuilder;
 import android.tools.traces.io.ResultWriter;
@@ -36,9 +37,6 @@
 import android.view.Choreographer;
 
 import androidx.test.filters.SmallTest;
-import androidx.test.platform.app.InstrumentationRegistry;
-
-import com.google.protobuf.InvalidProtocolBufferException;
 
 import org.junit.After;
 import org.junit.Before;
@@ -46,12 +44,9 @@
 import org.junit.Test;
 import org.mockito.Mockito;
 
-import perfetto.protos.PerfettoConfig.TracingServiceState;
 import perfetto.protos.PerfettoConfig.WindowManagerConfig.LogFrequency;
 
-import java.io.FileInputStream;
 import java.io.IOException;
-import java.util.Optional;
 
 /**
  * Test class for {@link WindowTracingPerfetto}.
@@ -74,7 +69,7 @@
         sChoreographer = Mockito.mock(Choreographer.class);
         sWindowTracing = new WindowTracingPerfetto(sWmMock, sChoreographer,
                 new WindowManagerGlobalLock(), TEST_DATA_SOURCE_NAME);
-        waitDataSourceIsAvailable();
+        busyWaitForDataSourceRegistration(TEST_DATA_SOURCE_NAME);
     }
 
     @Before
@@ -156,67 +151,4 @@
 
         mTraceMonitor.stop(writer);
     }
-
-    private static void waitDataSourceIsAvailable() {
-        final int timeoutMs = 10000;
-        final int busyWaitIntervalMs = 100;
-
-        int elapsedMs = 0;
-
-        while (!isDataSourceAvailable()) {
-            try {
-                Thread.sleep(busyWaitIntervalMs);
-                elapsedMs += busyWaitIntervalMs;
-                if (elapsedMs >= timeoutMs) {
-                    throw new RuntimeException("Data source didn't become available."
-                            + " Waited for: " + timeoutMs + " ms");
-                }
-            } catch (InterruptedException e) {
-                throw new RuntimeException(e);
-            }
-        }
-    }
-
-    private static boolean isDataSourceAvailable() {
-        byte[] proto = executeShellCommand("perfetto --query-raw");
-
-        try {
-            TracingServiceState state = TracingServiceState.parseFrom(proto);
-
-            Optional<Integer> producerId = Optional.empty();
-
-            for (TracingServiceState.Producer producer : state.getProducersList()) {
-                if (producer.getPid() == android.os.Process.myPid()) {
-                    producerId = Optional.of(producer.getId());
-                    break;
-                }
-            }
-
-            if (producerId.isEmpty()) {
-                return false;
-            }
-
-            for (TracingServiceState.DataSource ds : state.getDataSourcesList()) {
-                if (ds.getDsDescriptor().getName().equals(TEST_DATA_SOURCE_NAME)
-                        && ds.getProducerId() == producerId.get()) {
-                    return true;
-                }
-            }
-        } catch (InvalidProtocolBufferException e) {
-            throw new RuntimeException(e);
-        }
-
-        return false;
-    }
-
-    private static byte[] executeShellCommand(String command) {
-        try {
-            ParcelFileDescriptor fd = InstrumentationRegistry.getInstrumentation().getUiAutomation()
-                    .executeShellCommand(command);
-            FileInputStream is = new ParcelFileDescriptor.AutoCloseInputStream(fd);
-            return is.readAllBytes();
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
-    }
 }
diff --git a/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java b/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
index cfb2645..6f3deab 100644
--- a/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
+++ b/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
@@ -16,7 +16,7 @@
 
 package com.android.internal.protolog;
 
-import static android.tools.traces.Utils.executeShellCommand;
+import static android.tools.traces.Utils.busyWaitForDataSourceRegistration;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
@@ -50,7 +50,6 @@
 import com.android.internal.protolog.common.LogLevel;
 
 import com.google.common.truth.Truth;
-import com.google.protobuf.InvalidProtocolBufferException;
 
 import org.junit.After;
 import org.junit.Before;
@@ -60,14 +59,12 @@
 import org.junit.runners.JUnit4;
 import org.mockito.Mockito;
 
-import perfetto.protos.PerfettoConfig.TracingServiceState;
 import perfetto.protos.Protolog;
 import perfetto.protos.ProtologCommon;
 
 import java.io.File;
 import java.io.IOException;
 import java.util.List;
-import java.util.Optional;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -184,7 +181,7 @@
                     TestProtoLogGroup.values(), dataSourceBuilder, sProtoLogConfigurationService);
         }
 
-        waitDataSourceIsAvailable();
+        busyWaitForDataSourceRegistration(TEST_PROTOLOG_DATASOURCE_NAME);
     }
 
     @Before
@@ -870,54 +867,6 @@
                 .isEqualTo("This message should also be logged 567");
     }
 
-    private static void waitDataSourceIsAvailable() {
-        final int timeoutMs = 10000;
-        final int busyWaitIntervalMs = 100;
-
-        int elapsedMs = 0;
-
-        while (!isDataSourceAvailable()) {
-            SystemClock.sleep(busyWaitIntervalMs);
-            elapsedMs += busyWaitIntervalMs;
-            if (elapsedMs >= timeoutMs) {
-                throw new RuntimeException("Data source didn't become available."
-                        + " Waited for: " + timeoutMs + " ms");
-            }
-        }
-    }
-
-    private static boolean isDataSourceAvailable() {
-        byte[] proto = executeShellCommand("perfetto --query-raw");
-
-        try {
-            TracingServiceState state = TracingServiceState.parseFrom(proto);
-
-            Optional<Integer> producerId = Optional.empty();
-
-            for (TracingServiceState.Producer producer : state.getProducersList()) {
-                if (producer.getPid() == android.os.Process.myPid()) {
-                    producerId = Optional.of(producer.getId());
-                    break;
-                }
-            }
-
-            if (producerId.isEmpty()) {
-                return false;
-            }
-
-            for (TracingServiceState.DataSource ds : state.getDataSourcesList()) {
-                if (ds.getDsDescriptor().getName().equals(TEST_PROTOLOG_DATASOURCE_NAME)
-                        && ds.getProducerId() == producerId.get()) {
-                    return true;
-                }
-            }
-        } catch (InvalidProtocolBufferException e) {
-            throw new RuntimeException(e);
-        }
-
-        return false;
-    }
-
     private enum TestProtoLogGroup implements IProtoLogGroup {
         TEST_GROUP(true, true, false, "TEST_TAG");