FIX: Android can now start/stop while ios tests are running

maestro-cli/src/main/java/maestro/cli/command/TestCommand.kt

@@ -234,38 +234,71 @@ class TestCommand : Callable<Int> {
             host = parent?.host,
             port = parent?.port,
         ).map { it.instanceId }.toSet()
-        val deviceIds = getPassedOptionsDeviceIds()
+        val passedDeviceIds = getPassedOptionsDeviceIds()
             .filter { device ->
                 if (device !in availableDevices) {
                     throw CliError("Device $device was requested, but it is not connected.")
                 } else {
                     true
                 }
             }
-            .ifEmpty { availableDevices }
-            .toList()
-
-        val missingDevices = requestedShards - deviceIds.size
-        if (missingDevices > 0) {
-            PrintUtils.warn("Want to use ${deviceIds.size} devices, which is not enough to run $requestedShards shards. Missing $missingDevices device(s).")
-            throw CliError("Not enough devices connected ($missingDevices) to run the requested number of shards ($requestedShards).")
-        }
+            // Don't default to all available devices here. Handle it based on sharding later.
+            // .ifEmpty { availableDevices }
+            // .toList()
 
+        // Determine effective shards based on requested shards and flow count/type
         val effectiveShards = when {
+            onlySequenceFlows -> 1 // Sequence flows always run on one device
+            shardAll != null -> requestedShards.coerceAtMost(DeviceService.listConnectedDevices().size) // Cannot run more shards than connected devices if replicating all flows
+            shardSplit != null -> requestedShards.coerceAtMost(plan.flowsToRun.size).coerceAtMost(DeviceService.listConnectedDevices().size) // Cannot run more shards than flows or connected devices
+            else -> 1 // Default to 1 shard if no sharding options are provided
+        }
 
-            onlySequenceFlows -> 1
+        // Determine the list of device IDs to use for the shards
+        val deviceIds: List<String?> = if (passedDeviceIds.isNotEmpty()) {
+            // If specific devices were passed, use them
+            passedDeviceIds
+        } else {
+            // If no devices were passed...
+            if (effectiveShards == 1) {
+                // For a single shard, pass null to trigger device selection prompt if needed
+                listOf(null)
+            } else {
+                // For multiple shards, use all available connected devices.
+                // Note: This maintains existing behavior for sharding without explicit devices.
+                // The user might want prompting even for sharding, but that's a larger change.
+                val connectedDeviceIds = availableDevices.toList()
+                if (connectedDeviceIds.size < effectiveShards) {
+                    throw CliError("Not enough devices connected (${connectedDeviceIds.size}) to run the requested number of shards ($effectiveShards).")
+                }
+                // Use only as many devices as needed for the shards
+                connectedDeviceIds.take(effectiveShards)
+            }
+        }
 
-            shardAll == null -> requestedShards.coerceAtMost(plan.flowsToRun.size)
 
-            shardSplit == null -> requestedShards.coerceAtMost(deviceIds.size)
+        // Validate if enough devices are available for the requested shards IF specific devices were requested OR sharding is used without specific devices
+        if (passedDeviceIds.isNotEmpty() || (passedDeviceIds.isEmpty() && effectiveShards > 1)) {
+            val missingDevices = effectiveShards - deviceIds.size
+            if (missingDevices > 0) {
+                 PrintUtils.warn("Want to use ${deviceIds.size} devices, which is not enough to run $effectiveShards shards. Missing $missingDevices device(s).")
+                 throw CliError("Not enough devices available ($missingDevices missing) for the requested number of shards ($effectiveShards). Ensure devices are connected or reduce shard count.")
+            }
+        }
 
-            else -> 1
+        // Existing shard warning logic (adjusting message slightly)
+        if (shardAll == null && shardSplit != null && requestedShards > plan.flowsToRun.size) {
+            val warning = "Requested $requestedShards shards, " +
+                    "but cannot run more shards than flows (${plan.flowsToRun.size}). " +
+                    "Will use $effectiveShards shards instead."
+            PrintUtils.warn(warning)
+        } else if (shardAll != null && requestedShards > availableDevices.size) {
+             val warning = "Requested $requestedShards shards, " +
+                    "but cannot run more shards than connected devices (${availableDevices.size}). " +
+                    "Will use $effectiveShards shards instead."
+             PrintUtils.warn(warning)
         }
 
-        val warning = "Requested $requestedShards shards, " +
-                "but it cannot be higher than the number of flows (${plan.flowsToRun.size}). " +
-                "Will use $effectiveShards shards instead."
-        if (shardAll == null && requestedShards > plan.flowsToRun.size) PrintUtils.warn(warning)
 
         val chunkPlans = makeChunkPlans(plan, effectiveShards, onlySequenceFlows)
 
@@ -308,21 +341,28 @@ class TestCommand : Callable<Int> {
 
     private suspend fun runShardSuite(
         effectiveShards: Int,
-        deviceIds: List<String>,
+        deviceIds: List<String?>, // Allow null device IDs
         shardIndex: Int,
         chunkPlans: List<ExecutionPlan>,
         debugOutputPath: Path,
     ): Triple<Int?, Int?, TestExecutionSummary?> {
         val driverHostPort = selectPort(effectiveShards)
-        val deviceId = deviceIds[shardIndex]
+        // DeviceId can be null if we need to prompt
+        val deviceId: String? = deviceIds[shardIndex]
+
+        // Log device selection or indicate prompting might occur
+        if (deviceId != null) {
+            logger.info("[shard ${shardIndex + 1}] Selected device $deviceId using port $driverHostPort")
+        } else {
+            logger.info("[shard ${shardIndex + 1}] No device specified, will attempt connection or prompt for selection. Using port $driverHostPort")
+        }
 
-        logger.info("[shard ${shardIndex + 1}] Selected device $deviceId using port $driverHostPort")
 
         return MaestroSessionManager.newSession(
             host = parent?.host,
             port = parent?.port,
             driverHostPort = driverHostPort,
-            deviceId = deviceId,
+            deviceId = deviceId, // Pass the potentially null deviceId
             platform = parent?.platform,
             isHeadless = headless,
         ) { session ->

maestro-client/src/main/java/maestro/drivers/AndroidDriver.kt

@@ -70,15 +70,16 @@ class AndroidDriver(
     ) : Driver {
     private var open = false
     private val hostPort: Int = hostPort ?: DefaultDriverHostPort
+    // Track the actual port used (might differ from hostPort if there's a conflict)
+    private var actualPort: Int = hostPort ?: DefaultDriverHostPort
 
     private val metrics = metricsProvider.withPrefix("maestro.driver").withTags(mapOf("platform" to "android", "emulatorName" to emulatorName))
 
-    private val channel = ManagedChannelBuilder.forAddress("localhost", this.hostPort)
-        .usePlaintext()
-        .build()
-    private val blockingStub = MaestroDriverGrpc.newBlockingStub(channel)
+    // Initialize channel lazily after we know the actual port
+    private lateinit var channel: io.grpc.ManagedChannel
+    private val blockingStub get() = MaestroDriverGrpc.newBlockingStub(channel)
     private val blockingStubWithTimeout get() = blockingStub.withDeadlineAfter(120, TimeUnit.SECONDS)
-    private val asyncStub = MaestroDriverGrpc.newStub(channel)
+    private val asyncStub get() = MaestroDriverGrpc.newStub(channel)
     private val documentBuilderFactory = DocumentBuilderFactory.newInstance()
 
     private var instrumentationSession: AdbShellStream? = null
@@ -92,9 +93,20 @@ class AndroidDriver(
     }
 
     override fun open() {
+        // Allocate port forwarder first
         allocateForwarder()
+
+        // Get the actual port used (might have been changed during allocation)
+        actualPort = PORT_TO_FORWARDER.entries.find { it.value is AutoCloseable }?.key ?: hostPort
+
+        // Initialize the channel with the actual port
+        channel = ManagedChannelBuilder.forAddress("localhost", actualPort)
+            .usePlaintext()
+            .build()
+
+        // Install APKs and start instrumentation using the actual allocated port
         installMaestroApks()
-        startInstrumentationSession(hostPort)
+        startInstrumentationSession(actualPort)
 
         try {
             awaitLaunch()
@@ -141,30 +153,79 @@ class AndroidDriver(
 
 
     private fun allocateForwarder() {
+        // Close existing forwarder on this port if it exists
         PORT_TO_FORWARDER[hostPort]?.close()
-        PORT_TO_ALLOCATION_POINT[hostPort]?.let {
-            LOGGER.warn("Port $hostPort was already allocated. Allocation point: $it")
+
+        // Try to find an available port, starting with the preferred hostPort
+        val portToUse = findAvailablePort(hostPort)
+
+        if (portToUse != hostPort) {
+            LOGGER.info("Original port $hostPort was unavailable, using port $portToUse instead")
+        }
+
+        try {
+            PORT_TO_FORWARDER[portToUse] = dadb.tcpForward(
+                portToUse,
+                portToUse
+            )
+            PORT_TO_ALLOCATION_POINT[portToUse] = Exception().stackTraceToString()
+        } catch (e: Exception) {
+            LOGGER.error("Failed to allocate forwarder on port $portToUse: ${e.message}")
+            throw e
         }
-
-        PORT_TO_FORWARDER[hostPort] = dadb.tcpForward(
-            hostPort,
-            hostPort
-        )
-        PORT_TO_ALLOCATION_POINT[hostPort] = Exception().stackTraceToString()
+    }
+
+    /**
+     * Attempts to find an available port, starting with the preferred port.
+     * If that port is unavailable, tries ports in the range 7001-7128.
+     * 
+     * @param preferredPort The port to try first
+     * @return An available port or the original port if no alternatives found
+     */
+    private fun findAvailablePort(preferredPort: Int): Int {
+        // First try the preferred port
+        try {
+            // Quick check if we can bind to this port
+            val socket = java.net.ServerSocket(preferredPort)
+            socket.close()
+            return preferredPort
+        } catch (e: Exception) {
+            LOGGER.warn("Port $preferredPort is already in use, looking for alternative port")
+        }
+
+        // Try alternative ports in the range
+        for (port in (7001..7128).shuffled()) {
+            // Don't try the original port again
+            if (port == preferredPort) continue
+
+            try {
+                val socket = java.net.ServerSocket(port)
+                socket.close()
+                return port
+            } catch (e: Exception) {
+                // Port is in use, try next one
+                continue
+            }
+        }
+
+        // If we couldn't find an available port, return the original and let the
+        // operation fail, which will give a more descriptive error
+        LOGGER.warn("Could not find an available port in range 7001-7128, will try with original port $preferredPort")
+        return preferredPort
     }
 
     private fun awaitLaunch() {
         val startTime = System.currentTimeMillis()
 
         while (System.currentTimeMillis() - startTime < getStartupTimeout()) {
             runCatching {
-                dadb.open("tcp:$hostPort").close()
+                dadb.open("tcp:$actualPort").close()
                 return
             }
             Thread.sleep(100)
         }
 
-        throw AndroidDriverTimeoutException("Maestro Android driver did not start up in time  ---  emulator [ ${emulatorName} ] & port  [ dadb.open( tcp:${hostPort} ) ]")
+        throw AndroidDriverTimeoutException("Maestro Android driver did not start up in time  ---  emulator [ ${emulatorName} ] & port  [ dadb.open( tcp:${actualPort} ) ]")
     }
 
     override fun close() {
@@ -174,15 +235,15 @@ class AndroidDriver(
         }
 
         LOGGER.info("[Start] close port forwarder")
-        PORT_TO_FORWARDER[hostPort]?.close()
+        PORT_TO_FORWARDER[actualPort]?.close()
         LOGGER.info("[Done] close port forwarder")
 
         LOGGER.info("[Start] Remove host port from port forwarder map")
-        PORT_TO_FORWARDER.remove(hostPort)
+        PORT_TO_FORWARDER.remove(actualPort)
         LOGGER.info("[Done] Remove host port from port forwarder map")
 
         LOGGER.info("[Start] Remove host port from port to allocation map")
-        PORT_TO_ALLOCATION_POINT.remove(hostPort)
+        PORT_TO_ALLOCATION_POINT.remove(actualPort)
         LOGGER.info("[Done] Remove host port from port to allocation map")
 
         LOGGER.info("[Start] Uninstall driver from device")
@@ -204,6 +265,10 @@ class AndroidDriver(
     }
 
     override fun deviceInfo(): DeviceInfo {
+        if (!open) {
+            LOGGER.info("AndroidDriver.deviceInfo() called, but driver is not open. Calling open() first.")
+            open()
+        }
         return runDeviceCall {
             val response = blockingStubWithTimeout.deviceInfo(deviceInfoRequest {})