Make executeWorkbench work after VSCode restarts (#106)

* Make executeWorkbench work after VSCode restarts `vscode.openFolder` can be used to restart the VSCode under different workspace root. In that case the current socket is closed, and then the new connection appears after the VSCode is up again. This PR clears pending message callbacks when the connection is closed, and also re-assigns the `_promisedSocket` to a new socket when new connection appears. * Update test/specs/workbench.e2e.ts Co-authored-by: Sean Poulter <[email protected]> * Update test/specs/workbench.e2e.ts Co-authored-by: Sean Poulter <[email protected]> --------- Co-authored-by: Christian Bromann <[email protected]> Co-authored-by: Sean Poulter <[email protected]>
webdriverio-community · Feb 15, 2024 · 45de785 · 45de785
1 parent 0fc1795
commit 45de785
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/src/service.ts b/src/service.ts
@@ -55,6 +55,12 @@ export default class VSCodeWorkerService implements Services.ServiceInstance {
         }
     }
 
+    private _handleSocketClose (code: number, reason: Buffer) {
+        const msg = `Connection closed. Code: ${code}, reason: ${reason.toString()}`
+        this._promisedSocket = Promise.reject(new Error(msg))
+        this._pendingMessages.forEach((resolver) => resolver(msg, null))
+    }
+
     async beforeSession (option: Options.Testrunner, capabilities: VSCodeCapabilities) {
         this._isWebSession = capabilities.browserName !== 'vscode'
 
@@ -103,9 +109,11 @@ export default class VSCodeWorkerService implements Services.ServiceInstance {
                 )
                 wss.on('connection', (socket) => {
                     log.info('Connected with VSCode workbench')
+                    this._promisedSocket = Promise.resolve(socket)
                     resolve(socket)
                     clearTimeout(socketTimeout)
                     socket.on('message', this._handleIncoming.bind(this))
+                    socket.once('close', this._handleSocketClose.bind(this))
                 })
             })
         }

diff --git a/test/specs/workbench.e2e.ts b/test/specs/workbench.e2e.ts
@@ -3,6 +3,7 @@
 // eslint-disable-next-line @typescript-eslint/triple-slash-reference
 /// <reference path="../../dist/service.d.ts" />
 
+import path from 'node:path'
 import { browser, expect } from '@wdio/globals'
 import { skip } from './utils.js'
 
@@ -85,6 +86,20 @@ describe('workbench', () => {
         })
     })
 
+    it('can access the VSCode API after a new folder is opened @skipWeb', async () => {
+        const workspaceRoot = await browser.executeWorkbench<string | undefined>((vscode) => vscode.workspace.rootPath)
+        expect(workspaceRoot).not.toBe(undefined)
+        const newWorkspaceRoot = path.join(workspaceRoot!, 'test')
+        await browser.executeWorkbench((vscode, newRoot) => {
+            const uri = vscode.Uri.file(newRoot)
+            vscode.commands.executeCommand('vscode.openFolder', uri)
+        }, newWorkspaceRoot)
+        await browser.waitUntil(async () => {
+            const root = await browser.executeWorkbench((vscode) => vscode.workspace.rootPath)
+            return root === newWorkspaceRoot
+        })
+    })
+
     it('can send parameters to VSCode API invocation @skipWeb', async () => {
         const workbench = await browser.getWorkbench()
         const message = 'I passed this message as a parameter!'