Add getFutureForListener to TestUtils

Abstract out listener registration logic from getFutureForIntent to
apply to other listener registration callbacks.

Add tests which ensure that the listener is unregistered when the future
completes or is cancelled.

Bug: 288333346
Bug: 294636572
Test: atest mediatestutilstests
Change-Id: I91f34ee2e215cafe0d0cd0b33f831b244850ab06

2 files changed, 124 insertions, 46 deletions

diff --git a/media/tests/mediatestutils/java/com/android/media/mediatestutils/TestUtils.java b/media/tests/mediatestutils/java/com/android/media/mediatestutils/TestUtils.java
index 73dbe6780fec..69d836d24664 100644
--- a/media/tests/mediatestutils/java/com/android/media/mediatestutils/TestUtils.java
+++ b/media/tests/mediatestutils/java/com/android/media/mediatestutils/TestUtils.java
@@ -20,70 +20,100 @@ import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
-import android.util.Log;
 
 import androidx.concurrent.futures.CallbackToFutureAdapter;
 
-import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 
 import java.lang.ref.WeakReference;
-import java.util.concurrent.CancellationException;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeoutException;
-import java.util.concurrent.TimeUnit;
 import java.util.Objects;
+import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
-/**
- *
- */
+/** Utils for audio tests. */
 public class TestUtils {
-    public static final String TAG = "MediaTestUtils";
-
-    public static ListenableFuture<Intent> getFutureForIntent(Context context, String action,
-            Predicate<Intent> pred) {
+    /**
+     * Return a future for an intent delivered by a broadcast receiver which matches an
+     * action and predicate.
+     * @param context - Context to register the receiver with
+     * @param action - String representing action to register receiver for
+     * @param pred - Predicate which sets the future if evaluates to true, otherwise, leaves
+     * the future unset. If the predicate throws, the future is set exceptionally
+     * @return - The future representing intent delivery matching predicate.
+     */
+    public static ListenableFuture<Intent> getFutureForIntent(
+            Context context, String action, Predicate<Intent> pred) {
         // These are evaluated async
         Objects.requireNonNull(action);
         Objects.requireNonNull(pred);
-        // Doesn't need to be thread safe since the resolver is called inline
-        final WeakReference<BroadcastReceiver> wrapper[] = new WeakReference[1];
-        ListenableFuture<Intent> future = CallbackToFutureAdapter.getFuture(completer -> {
-            var receiver = new BroadcastReceiver() {
-                @Override
-                public void onReceive(Context context, Intent intent) {
+        return getFutureForListener(
+                (recv) ->
+                        context.registerReceiver(
+                                recv, new IntentFilter(action), Context.RECEIVER_NOT_EXPORTED),
+                (recv) -> {
                     try {
-                        if (action.equals(intent.getAction()) && pred.test(intent)) {
-                            completer.set(intent);
-                        }
-                    } catch (Exception e) {
-                        completer.setException(e);
+                        context.unregisterReceiver(recv);
+                    } catch (IllegalArgumentException e) {
+                        // Thrown when receiver is already unregistered, nothing to do
                     }
-                }
-            };
-            wrapper[0] = new WeakReference(receiver);
-            context.registerReceiver(receiver, new IntentFilter(action),
-                    Context.RECEIVER_NOT_EXPORTED);
-            return "Intent receiver future for ";
-        });
+                },
+                (completer) ->
+                        new BroadcastReceiver() {
+                            @Override
+                            public void onReceive(Context context, Intent intent) {
+                                try {
+                                    if (action.equals(intent.getAction()) && pred.test(intent)) {
+                                        completer.set(intent);
+                                    }
+                                } catch (Exception e) {
+                                    completer.setException(e);
+                                }
+                            }
+                        },
+                "Intent receiver future for action: " + action);
+    }
+
+    /**
+     * Return a future for a callback registered to a listener interface.
+     * @param registerFunc - Function which consumes the callback object for registration
+     * @param unregisterFunc - Function which consumes the callback object for unregistration
+     * This function is called when the future is completed or cancelled
+     * @param instantiateCallback - Factory function for the callback object, provided a completer
+     * object (see {@code CallbackToFutureAdapter.Completer<T>}), which is a logical reference
+     * to the future returned by this function
+     * @param debug - Debug string contained in future {@code toString} representation.
+     */
+    public static <T, V> ListenableFuture<T> getFutureForListener(
+            Consumer<V> registerFunc,
+            Consumer<V> unregisterFunc,
+            Function<CallbackToFutureAdapter.Completer<T>, V> instantiateCallback,
+            String debug) {
+        // Doesn't need to be thread safe since the resolver is called inline
+        final WeakReference<V> wrapper[] = new WeakReference[1];
+        ListenableFuture<T> future =
+                CallbackToFutureAdapter.getFuture(
+                        completer -> {
+                            final var cb = instantiateCallback.apply(completer);
+                            wrapper[0] = new WeakReference(cb);
+                            registerFunc.accept(cb);
+                            return debug;
+                        });
         if (wrapper[0] == null) {
-            throw new AssertionError("CallbackToFutureAdapter resolver should be called inline");
+            throw new AssertionError("Resolver should be called inline");
         }
         final var weakref = wrapper[0];
-        future.addListener(() -> {
-            try {
-                var recv = weakref.get();
-                // If there is no reference left, the receiver has already been unregistered
-                if (recv != null) {
-                    context.unregisterReceiver(recv);
-                    return;
-                }
-            } catch (IllegalArgumentException e) {
-                // Receiver already unregistered, nothing to do.
-            }
-            Log.d(TAG, "Intent receiver future for action: " + action +
-                    "unregistered prior to future completion/cancellation.");
-        } , MoreExecutors.directExecutor()); // Direct executor is fine since lightweight
+        future.addListener(
+                () -> {
+                    var cb = weakref.get();
+                    // If there is no reference left, the receiver has already been unregistered
+                    if (cb != null) {
+                        unregisterFunc.accept(cb);
+                        return;
+                    }
+                },
+                MoreExecutors.directExecutor()); // Direct executor is fine since lightweight
         return future;
     }
 }
diff --git a/media/tests/mediatestutils/tests/src/java/com/android/media/mediatestutils/GetFutureForIntentTest.java b/media/tests/mediatestutils/tests/src/java/com/android/media/mediatestutils/GetFutureForIntentTest.java
index 70b46ad86448..a4266a33f522 100644
--- a/media/tests/mediatestutils/tests/src/java/com/android/media/mediatestutils/GetFutureForIntentTest.java
+++ b/media/tests/mediatestutils/tests/src/java/com/android/media/mediatestutils/GetFutureForIntentTest.java
@@ -19,6 +19,7 @@ package com.android.media.mediatestutils;
 import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 
 import static com.android.media.mediatestutils.TestUtils.getFutureForIntent;
+import static com.android.media.mediatestutils.TestUtils.getFutureForListener;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -30,6 +31,8 @@ import android.os.SystemClock;
 
 import androidx.test.runner.AndroidJUnit4;
 
+import com.google.common.util.concurrent.ListenableFuture;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -100,6 +103,51 @@ public class GetFutureForIntentTest {
         assertThat(intent.getIntExtra(INTENT_EXTRA, -1)).isEqualTo(MAGIC_VALUE);
     }
 
+    @Test
+    public void unregisterListener_whenComplete() throws Exception {
+        final var service = new FakeService();
+        final ListenableFuture<Void> future =
+                getFutureForListener(
+                        service::registerListener,
+                        service::unregisterListener,
+                        completer ->
+                                () -> {
+                                    completer.set(null);
+                                },
+                        "FakeService listener future");
+        service.mRunnable.run();
+        assertThat(service.mRunnable).isNull();
+    }
+
+    @Test
+    public void unregisterListener_whenCancel() throws Exception {
+        final var service = new FakeService();
+        final ListenableFuture<Void> future =
+                getFutureForListener(
+                        service::registerListener,
+                        service::unregisterListener,
+                        completer ->
+                                () -> {
+                                    completer.set(null);
+                                },
+                        "FakeService listener future");
+        future.cancel(false);
+        assertThat(service.mRunnable).isNull();
+    }
+
+    private static class FakeService {
+        Runnable mRunnable;
+
+        void registerListener(Runnable r) {
+            mRunnable = r;
+        }
+
+        void unregisterListener(Runnable r) {
+            assertThat(r).isEqualTo(mRunnable);
+            mRunnable = null;
+        }
+    }
+
     private void sendIntent(boolean correctValue) {
         final Intent intent = new Intent(INTENT_ACTION).setPackage(mContext.getPackageName());
         intent.putExtra(INTENT_EXTRA, correctValue ? MAGIC_VALUE : MAGIC_VALUE + 1);