Fix teardown and tests

MetaMask · Nov 9, 2023 · f5fe7f6 · f5fe7f6
1 parent 0e044db
commit f5fe7f6
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 106 deletions.
diff --git a/packages/snaps-execution-environments/coverage.json b/packages/snaps-execution-environments/coverage.json
@@ -1,6 +1,6 @@
 {
-  "branches": 80.71,
-  "functions": 92.02,
-  "lines": 91.66,
-  "statements": 91.36
+  "branches": 80.28,
+  "functions": 90.41,
+  "lines": 91.06,
+  "statements": 90.65
 }
diff --git a/packages/snaps-execution-environments/src/common/BaseSnapExecutor.test.browser.ts b/packages/snaps-execution-environments/src/common/BaseSnapExecutor.test.browser.ts
@@ -1685,6 +1685,7 @@ describe('BaseSnapExecutor', () => {
     expect(await executor.readCommand()).toStrictEqual({
       jsonrpc: '2.0',
       method: 'OutboundResponse',
+      params: { source: 'ethereum.request' },
     });
 
     expect(await executor.readCommand()).toStrictEqual({
@@ -1796,6 +1797,7 @@ describe('BaseSnapExecutor', () => {
     expect(await executor.readCommand()).toStrictEqual({
       jsonrpc: '2.0',
       method: 'OutboundResponse',
+      params: { source: 'ethereum.request' },
     });
 
     expect(await executor.readCommand()).toStrictEqual({

diff --git a/packages/snaps-execution-environments/src/common/endowments/network.test.ts b/packages/snaps-execution-environments/src/common/endowments/network.test.ts
@@ -53,21 +53,14 @@ describe('Network endowments', () => {
     });
 
     it('can use AbortController normally', async () => {
-      let resolve: ((result: string) => void) | null = null;
-      fetchMock.mockOnce(
-        async () =>
-          new Promise((_resolve) => {
-            resolve = _resolve;
-          }),
-      );
+      fetchMock.mockOnce(async () => 'FAIL');
 
       const { fetch } = network.factory(factoryOptions);
 
       const controller = new AbortController();
 
       const fetchPromise = fetch('foo.com', { signal: controller.signal });
       controller.abort();
-      (resolve as any)('FAIL');
       await expect(fetchPromise).rejects.toThrow('The operation was aborted');
     });
 
@@ -86,10 +79,7 @@ describe('Network endowments', () => {
     it.todo('reason from AbortController.abort() is passed down');
 
     it('should not expose then or catch after teardown has been called', async () => {
-      let fetchResolve: ((result: string) => void) | null = null;
-      fetchMock.mockOnce(
-        async () => new Promise((resolve) => (fetchResolve = resolve)),
-      );
+      fetchMock.mockOnce(async () => 'Resolved');
 
       const { fetch, teardownFunction } = network.factory(factoryOptions);
       const ErrorProxy = jest
@@ -107,7 +97,6 @@ describe('Network endowments', () => {
         .catch((error) => console.log(error));
 
       const teardownPromise = teardownFunction();
-      (fetchResolve as any)('Resolved');
       await teardownPromise;
       await new Promise((resolve) => setTimeout(() => resolve('Resolved'), 0));
 

diff --git a/packages/snaps-execution-environments/src/common/endowments/network.ts b/packages/snaps-execution-environments/src/common/endowments/network.ts
@@ -65,33 +65,45 @@ class ResponseWrapper implements Response {
   }
 
   async text() {
-    await this.#onStart();
-    try {
-      return await withTeardown<string>(this.#ogResponse.text(), this as any);
-    } finally {
-      await this.#onFinish();
-    }
+    return await withTeardown<string>(
+      (async () => {
+        await this.#onStart();
+        try {
+          return await this.#ogResponse.text();
+        } finally {
+          await this.#onFinish();
+        }
+      })(),
+      this.#teardownRef,
+    );
   }
 
   async arrayBuffer(): Promise<ArrayBuffer> {
-    await this.#onStart();
-    try {
-      return await withTeardown<ArrayBuffer>(
-        this.#ogResponse.arrayBuffer(),
-        this as any,
-      );
-    } finally {
-      await this.#onFinish();
-    }
+    return await withTeardown<ArrayBuffer>(
+      (async () => {
+        await this.#onStart();
+        try {
+          return await this.#ogResponse.arrayBuffer();
+        } finally {
+          await this.#onFinish();
+        }
+      })(),
+      this.#teardownRef,
+    );
   }
 
   async blob(): Promise<Blob> {
-    await this.#onStart();
-    try {
-      return await withTeardown<Blob>(this.#ogResponse.blob(), this as any);
-    } finally {
-      await this.#onFinish();
-    }
+    return await withTeardown<Blob>(
+      (async () => {
+        await this.#onStart();
+        try {
+          return await this.#ogResponse.blob();
+        } finally {
+          await this.#onFinish();
+        }
+      })(),
+      this.#teardownRef,
+    );
   }
 
   clone(): Response {
@@ -105,24 +117,31 @@ class ResponseWrapper implements Response {
   }
 
   async formData(): Promise<FormData> {
-    await this.#onStart();
-    try {
-      return await withTeardown<FormData>(
-        this.#ogResponse.formData(),
-        this as any,
-      );
-    } finally {
-      await this.#onFinish();
-    }
+    return await withTeardown<FormData>(
+      (async () => {
+        await this.#onStart();
+        try {
+          return await this.#ogResponse.formData();
+        } finally {
+          await this.#onFinish();
+        }
+      })(),
+      this.#teardownRef,
+    );
   }
 
   async json(): Promise<any> {
-    await this.#onStart();
-    try {
-      return await withTeardown(this.#ogResponse.json(), this as any);
-    } finally {
-      await this.#onFinish();
-    }
+    return await withTeardown(
+      (async () => {
+        await this.#onStart();
+        try {
+          return await this.#ogResponse.json();
+        } finally {
+          await this.#onFinish();
+        }
+      })(),
+      this.#teardownRef,
+    );
   }
 }
 
@@ -195,60 +214,71 @@ const createNetwork = ({ notify }: EndowmentFactoryOptions = {}) => {
 
     let res: Response;
     let openFetchConnection: { cancel: () => Promise<void> } | undefined;
-    try {
-      await notify({ method: 'OutboundRequest', params: { source: 'fetch' } });
-      const fetchPromise = fetch(input, {
-        ...init,
-        signal: abortController.signal,
-      });
-
-      openFetchConnection = {
-        cancel: async () => {
-          abortController.abort();
-          try {
-            await fetchPromise;
-          } catch {
-            /* do nothing */
+    return await withTeardown(
+      (async () => {
+        try {
+          await notify({
+            method: 'OutboundRequest',
+            params: { source: 'fetch' },
+          });
+          const fetchPromise = fetch(input, {
+            ...init,
+            signal: abortController.signal,
+          });
+
+          openFetchConnection = {
+            cancel: async () => {
+              abortController.abort();
+              try {
+                await fetchPromise;
+              } catch {
+                /* do nothing */
+              }
+            },
+          };
+          openConnections.add(openFetchConnection);
+
+          res = new ResponseWrapper(
+            await fetchPromise,
+            teardownRef,
+            onStart,
+            onFinish,
+          );
+        } finally {
+          if (openFetchConnection !== undefined) {
+            openConnections.delete(openFetchConnection);
           }
-        },
-      };
-      openConnections.add(openFetchConnection);
-
-      res = new ResponseWrapper(
-        await withTeardown(fetchPromise, teardownRef),
-        teardownRef,
-        onStart,
-        onFinish,
-      );
-    } finally {
-      if (openFetchConnection !== undefined) {
-        openConnections.delete(openFetchConnection);
-      }
-      await notify({ method: 'OutboundResponse', params: { source: 'fetch' } });
-    }
-
-    if (res.body !== null) {
-      const body = new WeakRef<ReadableStream>(res.body);
-
-      const openBodyConnection = {
-        cancel:
-          /* istanbul ignore next: see it.todo('can be torn down during body read') test */
-          async () => {
-            try {
-              await body.deref()?.cancel();
-            } catch {
-              /* do nothing */
-            }
-          },
-      };
-      openConnections.add(openBodyConnection);
-      cleanup.register(
-        res.body,
-        /* istanbul ignore next: can't test garbage collection without modifying node parameters */
-        () => openConnections.delete(openBodyConnection),
-      );
-    }
-    return harden(res);
+          await notify({
+            method: 'OutboundResponse',
+            params: { source: 'fetch' },
+          });
+        }
+
+        if (res.body !== null) {
+          const body = new WeakRef<ReadableStream>(res.body);
+
+          const openBodyConnection = {
+            cancel:
+              /* istanbul ignore next: see it.todo('can be torn down during body read') test */
+              async () => {
+                try {
+                  await body.deref()?.cancel();
+                } catch {
+                  /* do nothing */
+                }
+              },
+          };
+          openConnections.add(openBodyConnection);
+          cleanup.register(
+            res.body,
+            /* istanbul ignore next: can't test garbage collection without modifying node parameters */
+            () => openConnections.delete(openBodyConnection),
+          );
+        }
+        return harden(res);
+      })(),
+      teardownRef,
+    );
   };
 
   const teardownFunction = async () => {