Simplify WebViewMessageStream

Author: FrederikBolding

packages/snaps-controllers/src/services/webview/WebViewExecutionService.ts

@@ -45,7 +45,7 @@ export class WebViewExecutionService extends AbstractExecutionService<string> {
     const stream = new WebViewMessageStream({
       name: 'parent',
       target: 'child',
-      getWebView: async () => webView,
+      webView,
     });
 
     return { worker: jobId, stream };

packages/snaps-controllers/src/services/webview/WebViewMessageStream.ts

@@ -1,7 +1,6 @@
 import type { PostMessageEvent } from '@metamask/post-message-stream';
 import { BasePostMessageStream } from '@metamask/post-message-stream';
 import { isValidStreamMessage } from '@metamask/post-message-stream/dist/utils';
-import { logError } from '@metamask/snaps-utils';
 import { assert, stringToBytes } from '@metamask/utils';
 
 export type WebViewInterface = {
@@ -13,7 +12,7 @@ export type WebViewInterface = {
 export type WebViewStreamArgs = {
   name: string;
   target: string;
-  getWebView: () => Promise<WebViewInterface>;
+  webView: WebViewInterface;
 };
 
 /**
@@ -33,29 +32,21 @@ export class WebViewMessageStream extends BasePostMessageStream {
    * @param args.name - The name of the stream. Used to differentiate between
    * multiple streams sharing the same window object.
    * @param args.target - The name of the stream to exchange messages with.
-   * @param args.getWebView - A asynchronous getter for the webview.
+   * @param args.webView - A reference to the WebView.
    */
-  constructor({ name, target, getWebView }: WebViewStreamArgs) {
+  constructor({ name, target, webView }: WebViewStreamArgs) {
     super();
 
     this.#name = name;
     this.#target = target;
 
     this._onMessage = this._onMessage.bind(this);
 
-    // This is a bit atypical from other post-message streams.
-    // We have to wait for the WebView to fully load before we can continue using the stream.
-    getWebView()
-      .then((webView) => {
-        this.#webView = webView;
-        // This method is already bound.
-        // eslint-disable-next-line @typescript-eslint/unbound-method
-        webView.registerMessageListener(this._onMessage);
-        this._handshake();
-      })
-      .catch((error) => {
-        logError(error);
-      });
+    this.#webView = webView;
+    // This method is already bound.
+    // eslint-disable-next-line @typescript-eslint/unbound-method
+    this.#webView.registerMessageListener(this._onMessage);
+    this._handshake();
   }
 
   protected _postMessage(data: unknown): void {

packages/snaps-controllers/src/test-utils/webview.ts

@@ -17,7 +17,7 @@ export function parseInjectedJS(js: string) {
 /**
  * Takes no param and return mocks necessary for testing WebViewMessageStream.
  *
- * @returns The mockWebView, mockGetWebView, and mockStream.
+ * @returns The mockWebView, and mockStream.
  */
 export function createWebViewObjects() {
   const registerMessageListenerA = jest.fn();
@@ -45,26 +45,21 @@ export function createWebViewObjects() {
     }),
   };
 
-  const mockGetWebViewA = jest.fn().mockResolvedValue(mockWebViewA);
-  const mockGetWebViewB = jest.fn().mockResolvedValue(mockWebViewB);
-
   const streamA = new WebViewMessageStream({
     name: 'a',
     target: 'b',
-    getWebView: mockGetWebViewA,
+    webView: mockWebViewA,
   });
 
   const streamB = new WebViewMessageStream({
     name: 'b',
     target: 'a',
-    getWebView: mockGetWebViewB,
+    webView: mockWebViewB,
   });
 
   return {
     mockWebViewA,
     mockWebViewB,
-    mockGetWebViewA,
-    mockGetWebViewB,
     streamA,
     streamB,
   };