register websocket when appium throws unexpectedError (#1059)

AppiumTestDistribution · Mar 20, 2024 · 1de3570 · 1de3570
1 parent dc242e2
commit 1de3570
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/src/plugin.ts b/src/plugin.ts
@@ -69,7 +69,7 @@ import * as https from 'https';
 import fs from 'fs';
 // import { runAndroidAdbServer } from './device-managers/AdbServer';
 import { getStreamingServer } from './streaming-server';
-import { sleep } from 'asyncbox';
+import { sleep, waitForCondition } from 'asyncbox';
 
 const commandsQueueGuard = new AsyncLock();
 const DEVICE_MANAGER_LOCK_NAME = 'DeviceManager';
@@ -117,6 +117,13 @@ class DevicePlugin extends BasePlugin {
         deviceFilter,
       )} onUnexpectedShutdown from server`,
     );
+    await waitForCondition(async () => {
+      console.log('Waiting for websocket handlers to be removed..');
+      return !(await DevicePlugin.httpServer.getWebSocketHandlers()).hasOwnProperty(
+        '/android-stream/:udid',
+      );
+    });
+    await DevicePlugin.registerWebSocketHandlers();
   }
 
   private static async registerWebSocketHandlers() {
@@ -249,6 +256,7 @@ class DevicePlugin extends BasePlugin {
               try {
                 const selectedDevice = client.getDevice(device.id);
                 await selectedDevice.uninstall('com.device.farm');
+                await sleep(800);
                 await selectedDevice.install(apk);
                 console.log(`Installed ${apk} on device ${device.id}`);
               } catch (installError: any) {

diff --git a/src/streaming-server.ts b/src/streaming-server.ts
@@ -84,7 +84,7 @@ export function getStreamingServer(adbInstance: ADB) {
         await proxySocket.start();
         log.info(`Successfully created proxy streaming session for device ${udid}`);
         SocketMap.set(streamingKey, proxySocket);
-        proxySocket.addListener(ws);
+        await proxySocket.addListener(ws);
       } catch (err) {
         console.log('Unable to connect to remote streaming socket');
         console.log(err);

diff --git a/web/src/components/AndroidStream.tsx b/web/src/components/AndroidStream.tsx
@@ -12,7 +12,10 @@ const AndroidStream = () => {
   const createWebSocketConnection = (wsUrl: string) => {
     ws = new WebSocket(wsUrl);
     ws.addEventListener('message', handleWebSocketMessage);
-    ws.addEventListener('close', () => createWebSocketConnection(wsUrl));
+    ws.addEventListener('close', () => {
+      console.log('WebSocket connection closed. Reconnecting...');
+      createWebSocketConnection(wsUrl)
+    });
   };
 
   useEffect(() => {