diff --git a/packages/snaps-controllers/src/websocket/WebSocketService.test.ts b/packages/snaps-controllers/src/websocket/WebSocketService.test.ts index 4664a8d8c8..3cfff491b0 100644 --- a/packages/snaps-controllers/src/websocket/WebSocketService.test.ts +++ b/packages/snaps-controllers/src/websocket/WebSocketService.test.ts @@ -32,8 +32,10 @@ class MockWebSocket { } else if (type === 'close') { this.#closeListener = listener; } else if (type === 'error' && this.#origin === MOCK_WEBSOCKET_BROKEN_URI) { - listener(new Event('error')); - this.#closeWithCode(1006, '', false); + setTimeout(() => { + listener(new Event('error')); + this.#closeWithCode(1006, '', false); + }, 1); } else if (type === 'message') { this.#messageListener = listener; } @@ -259,6 +261,10 @@ describe('WebSocketService', () => { MOCK_WEBSOCKET_BROKEN_URI, ), ).rejects.toThrow('An error occurred while opening the WebSocket.'); + + expect( + messenger.call('WebSocketService:getAll', MOCK_SNAP_ID), + ).toHaveLength(0); }); it('logs if the Snap request fails', async () => { diff --git a/packages/snaps-controllers/src/websocket/WebSocketService.ts b/packages/snaps-controllers/src/websocket/WebSocketService.ts index 17a6ba29e5..3fcc760036 100644 --- a/packages/snaps-controllers/src/websocket/WebSocketService.ts +++ b/packages/snaps-controllers/src/websocket/WebSocketService.ts @@ -219,6 +219,8 @@ export class WebSocketService { }); socket.addEventListener('close', (event) => { + this.#sockets.delete(id); + this.#handleEvent(snapId, { type: 'close', id, @@ -284,8 +286,6 @@ export class WebSocketService { const { socket } = this.#get(snapId, id); socket.close(); - - this.#sockets.delete(id); } /**