Merge "Make sure sendresult after association was created" into sc-v2-dev
diff --git a/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl b/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl
index a630873..e1c13f7 100644
--- a/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl
+++ b/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl
@@ -29,4 +29,6 @@
         in String callingPackage,
         in IFindDeviceCallback findCallback,
         in AndroidFuture<Association> serviceCallback);
+
+    void onAssociationCreated();
 }
diff --git a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceActivity.java b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceActivity.java
index b4cafd8..edbfd2a 100644
--- a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceActivity.java
+++ b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceActivity.java
@@ -254,8 +254,14 @@
         Log.i(LOG_TAG, "onDeviceConfirmed(selectedDevice = " + selectedDevice + ")");
         getService().onDeviceSelected(
                 getCallingPackage(), getDeviceMacAddress(selectedDevice.device));
+    }
+
+    void setResultAndFinish() {
+        Log.i(LOG_TAG, "setResultAndFinish(selectedDevice = "
+                + getService().mSelectedDevice.device + ")");
         setResult(RESULT_OK,
-                new Intent().putExtra(CompanionDeviceManager.EXTRA_DEVICE, selectedDevice.device));
+                new Intent().putExtra(
+                        CompanionDeviceManager.EXTRA_DEVICE, getService().mSelectedDevice.device));
         finish();
     }
 
diff --git a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java
index c24782e..5df8e3c 100644
--- a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java
+++ b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java
@@ -117,6 +117,11 @@
                     CompanionDeviceDiscoveryService::startDiscovery,
                     CompanionDeviceDiscoveryService.this, request));
         }
+
+        @Override
+        public void onAssociationCreated() {
+            Handler.getMain().post(CompanionDeviceDiscoveryService.this::onAssociationCreated);
+        }
     };
 
     private ScanCallback mBLEScanCallback;
@@ -222,6 +227,11 @@
                 SCAN_TIMEOUT);
     }
 
+    @MainThread
+    private void onAssociationCreated() {
+        mActivity.setResultAndFinish();
+    }
+
     private boolean shouldScan(List<? extends DeviceFilter> mediumSpecificFilters) {
         return !isEmpty(mediumSpecificFilters) || isEmpty(mFilters);
     }
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 4748a86..2028766 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -457,6 +457,9 @@
             }, FgThread.getExecutor()).whenComplete(uncheckExceptions((association, err) -> {
                 if (err == null) {
                     addAssociation(association, userId);
+                    mServiceConnectors.forUser(userId).post(service -> {
+                        service.onAssociationCreated();
+                    });
                 } else {
                     Slog.e(LOG_TAG, "Failed to discover device(s)", err);
                     callback.onFailure("No devices found: " + err.getMessage());