From d6bc119f8815074f34e8487574dc89c696764e0f Mon Sep 17 00:00:00 2001
From: JiaLiPassion <JiaLi.Passion@gmail.com>
Date: Thu, 12 Mar 2020 00:50:47 +0900
Subject: [PATCH] fix: zone.js patch jest should handle done correctly

in the pervious PR #340, zone.js add support to jest test.each methods, but it
introduces a bug, which is the `done()` function will not be handled correctly.

```
it('should work with done', done => {
  // done will be undefined.
});
```

The reason is the logic of monkey patching `test` method changed from

```
return testBody.length === 0
   ? () => testProxyZone.run(testBody, null)
   : done => testProxyZone.run(testBody, null, [done]);
```
to
```
 return function(...args) {
   return testProxyZone.run(testBody, null, args);
 };
```

the purpose of this change is to handle the following cases.

```
test.each([1, 2])('test.each', (arg1, arg2) => {
  expect(arg1).toBe(1);
  expect(arg2).toBe(2);
});
```

so in jest, it is a little complex, because the `testBody`'s parameter may be bigger than 1, so the
previous logic

```
return testBody.length === 0
   ? () => testProxyZone.run(testBody, null)
   : done => testProxyZone.run(testBody, null, [done]);
```
will not work for `test.each`.

So in this PR, I created a dynamic `Function` to return the correct length of paramters (which is required by jest core), to handle
1. normal `test` with or without `done`.
2. each with parameters with or without done.
---
 example/src/app/app.component.spec.ts | 20 +++++++++
 src/zone-patch/index.js               | 58 ++++++++++++++++++++++++---
 2 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/example/src/app/app.component.spec.ts b/example/src/app/app.component.spec.ts
index 910f8867ba..46903c06ca 100644
--- a/example/src/app/app.component.spec.ts
+++ b/example/src/app/app.component.spec.ts
@@ -79,6 +79,26 @@ describe('AppComponent', () => {
       done();
     }, 100);
   });
+
+  it('async with done should work', async done => {
+    let flag = false;
+    setTimeout(() => {
+      flag = true;
+      expect(flag).toBe(true);
+      done();
+    }, 100);
+  });
+
+  it.each([[1, 2]])('it.each', (arg1, arg2) => {
+    expect(arg1).toBe(1);
+    expect(arg2).toBe(2);
+  });
+
+  (it.each([[1, 2]]) as any)('it.each with done', (arg1, arg2, done) => {
+    expect(arg1).toBe(1);
+    expect(arg2).toBe(2);
+    done();
+  });
 });
 
 test.todo('a sample todo');
diff --git a/src/zone-patch/index.js b/src/zone-patch/index.js
index 8c2631e1dc..25a0d4a39a 100644
--- a/src/zone-patch/index.js
+++ b/src/zone-patch/index.js
@@ -44,13 +44,51 @@ function wrapDescribeInZone(describeBody) {
 // Create a proxy zone in which to run `test` blocks so that the tests function
 // can retroactively install different zones.
 const testProxyZone = ambientZone.fork(new ProxyZoneSpec());
-function wrapTestInZone(testBody) {
+function wrapTestInZone(testBody, eachArgs) {
   if (testBody === undefined) {
     return;
   }
-  return function(...args) {
-    return testProxyZone.run(testBody, null, args);
-  };
+
+  if (!eachArgs || eachArgs.length === 0 || eachArgs[0].length === 0) {
+    // If we are not handling `test.each`, then the parameter of `testBody`
+    // will be 0 or 1, if it is 1, then we need to return a function with
+    // done parameter
+    return testBody.length === 0
+      ? () => testProxyZone.run(testBody, null)
+      : done => testProxyZone.run(testBody, null, [done]);
+  } else {
+    // Dynamically create a Function to contain the same length
+    // of the parameters as the testBody
+    // For example:
+    // ```
+    // test.each([[1, 2]])('test.each', (arg1, arg2) => {});
+    // ```
+    // In this case we need to return a function like this
+    // ```
+    // return function(arg1, arg2) {
+    //   return testProxyZone.run(testBody, null, [arg1, arg2]);
+    // }
+    // ```
+    const len = eachArgs[0].length;
+    const args = [];
+    let argString = '';
+    for (let i = 0; i < len; i++) {
+      args.push('arg' + i);
+      argString += 'arg' + i;
+      if (i !== len - 1) {
+        argString += ', ';
+      }
+    }
+    args.push('testBody');
+    args.push('testProxyZone');
+    if (len < testBody.length) {
+      args.push('done');
+      argString += ', done';
+    }
+    const funcBody = `
+      return testProxyZone.run(testBody, null, [${argString}])`;
+    return new Function(args, funcBody);
+  }
 }
 
 /**
@@ -70,7 +108,17 @@ const bindDescribe = originalJestFn =>
 const bindTest = originalJestFn =>
   function(...eachArgs) {
     return function(...args) {
-      args[1] = wrapTestInZone(args[1]);
+      const testBody = args[1];
+      args[1] = wrapTestInZone(args[1], ...eachArgs);
+      if (testBody.length > 0 || (eachArgs.length > 0 && eachArgs[0].length > 0)) {
+        eachArgs.forEach(row => {
+          const modifiedRow = row.map(a => {
+            a.push(testBody);
+            a.push(testProxyZone);
+          });
+          return modifiedRow;
+        });
+      }
       return originalJestFn.apply(this, eachArgs).apply(this, args);
     };
   };