From d38baae01c4a6848e696090ba671aa442912f6ed Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Thu, 2 Apr 2020 13:12:22 -0700
Subject: [PATCH] feat(testrunner): nested environments (#1635)

---
 utils/testrunner/TestRunner.js           | 39 ++++++++++++++++++------
 utils/testrunner/test/testrunner.spec.js | 33 ++++++++++++++------
 2 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/utils/testrunner/TestRunner.js b/utils/testrunner/TestRunner.js
index 9335b5dd1b604..fe4eaaee9b269 100644
--- a/utils/testrunner/TestRunner.js
+++ b/utils/testrunner/TestRunner.js
@@ -156,10 +156,20 @@ class Test {
   }
 
   environment(environment) {
+    const parents = new Set();
+    for (let parent = environment; !(parent instanceof Suite); parent = parent.parentEnvironment())
+      parents.add(parent);
+    for (const env of this._environments) {
+      for (let parent = env; !(parent instanceof Suite); parent = parent.parentEnvironment()) {
+        if (parents.has(parent))
+          throw new Error(`Cannot use environments "${environment.name()}" and "${env.name()}" that share a parent environment "${parent.fullName()}" in test "${this.fullName()}"`);
+      }
+    }
+    const environmentParentSuite = environment.parentSuite();
     for (let suite = this.suite(); suite; suite = suite.parentSuite()) {
-      if (suite === environment.parentSuite()) {
+      if (suite === environmentParentSuite) {
         this._environments.push(environment);
-        return;
+        return this;
       }
     }
     throw new Error(`Cannot use environment "${environment.name()}" from suite "${environment.parentSuite().fullName()}" in unrelated test "${this.fullName()}"`);
@@ -167,14 +177,21 @@ class Test {
 }
 
 class Environment {
-  constructor(parentSuite, name, location) {
-    this._parentSuite = parentSuite;
+  constructor(parentEnvironment, name, location) {
+    this._parentEnvironment = parentEnvironment;
+    this._parentSuite = parentEnvironment;
+    if (parentEnvironment && !(parentEnvironment instanceof Suite))
+      this._parentSuite = parentEnvironment.parentSuite();
     this._name = name;
-    this._fullName = (parentSuite ? parentSuite.fullName() + ' ' + name : name).trim();
+    this._fullName = (parentEnvironment ? parentEnvironment.fullName() + ' ' + name : name).trim();
     this._location = location;
     this._hooks = [];
   }
 
+  parentEnvironment() {
+    return this._parentEnvironment;
+  }
+
   parentSuite() {
     return this._parentSuite;
   }
@@ -397,8 +414,12 @@ class TestWorker {
     for (let suite = test.suite(); suite; suite = suite.parentSuite())
       environmentStack.push(suite);
     environmentStack.reverse();
-    for (const environment of test._environments)
-      environmentStack.splice(environmentStack.indexOf(environment.parentSuite()) + 1, 0, environment);
+    for (const environment of test._environments) {
+      const insert = [];
+      for (let parent = environment; !(parent instanceof Suite); parent = parent.parentEnvironment())
+        insert.push(parent);
+      environmentStack.splice(environmentStack.indexOf(environment.parentSuite()) + 1, 0, ...insert.reverse());
+    }
 
     let common = 0;
     while (common < environmentStack.length && this._environmentStack[common] === environmentStack[common])
@@ -677,13 +698,11 @@ class TestRunner extends EventEmitter {
     this.describe = this._suiteBuilder([]);
     this.it = this._testBuilder([]);
     this.environment = (name, callback) => {
-      if (!(this._currentEnvironment instanceof Suite))
-        throw new Error(`Cannot define an environment inside an environment`);
       const location = Location.getCallerLocation(__filename);
       const environment = new Environment(this._currentEnvironment, name, location);
       this._currentEnvironment = environment;
       callback();
-      this._currentEnvironment = environment.parentSuite();
+      this._currentEnvironment = environment.parentEnvironment();
       return environment;
     };
     this.Expectations = { ...TestExpectation };
diff --git a/utils/testrunner/test/testrunner.spec.js b/utils/testrunner/test/testrunner.spec.js
index 980463dc67213..f07acddb641d0 100644
--- a/utils/testrunner/test/testrunner.spec.js
+++ b/utils/testrunner/test/testrunner.spec.js
@@ -244,11 +244,16 @@ module.exports.addTests = function({testRunner, expect}) {
     it('should run all hooks in proper order', async() => {
       const log = [];
       const t = newTestRunner();
-      const e = t.environment('env', () => {
+      let e2;
+      t.environment('env', () => {
         t.beforeAll(() => log.push('env:beforeAll'));
         t.afterAll(() => log.push('env:afterAll'));
         t.beforeEach(() => log.push('env:beforeEach'));
         t.afterEach(() => log.push('env:afterEach'));
+        e2 = t.environment('env2', () => {
+          t.beforeAll(() => log.push('env2:beforeAll'));
+          t.afterAll(() => log.push('env2:afterAll'));
+        });
       });
       t.beforeAll(() => log.push('root:beforeAll'));
       t.beforeEach(() => log.push('root:beforeEach1'));
@@ -270,13 +275,13 @@ module.exports.addTests = function({testRunner, expect}) {
         t.afterEach(() => log.push('suite:afterEach2'));
         t.afterAll(() => log.push('suite:afterAll'));
       });
-      t.it('cuatro', () => log.push('test #4')).environment(e);
+      t.it('cuatro', () => log.push('test #4')).environment(e2);
       t.describe('no hooks suite', () => {
         t.describe('suite2', () => {
           t.beforeAll(() => log.push('suite2:beforeAll'));
           t.afterAll(() => log.push('suite2:afterAll'));
           t.describe('no hooks suite 2', () => {
-            t.it('cinco', () => log.push('test #5')).environment(e);
+            t.it('cinco', () => log.push('test #5')).environment(e2);
           });
         });
       });
@@ -321,6 +326,7 @@ module.exports.addTests = function({testRunner, expect}) {
         'suite:afterAll',
 
         'env:beforeAll',
+        'env2:beforeAll',
 
         'root:beforeEach1',
         'root:beforeEach2',
@@ -338,6 +344,7 @@ module.exports.addTests = function({testRunner, expect}) {
         'root:afterEach',
         'suite2:afterAll',
 
+        'env2:afterAll',
         'env:afterAll',
 
         'root:afterAll1',
@@ -347,8 +354,10 @@ module.exports.addTests = function({testRunner, expect}) {
     it('environment restrictions', async () => {
       const t = newTestRunner();
       let env;
+      let env2;
       t.describe('suite1', () => {
         env = t.environment('env', () => {
+          env2 = t.environment('env2', () => {});
           try {
             t.it('test', () => {});
             expect(true).toBe(false);
@@ -361,13 +370,19 @@ module.exports.addTests = function({testRunner, expect}) {
           } catch (e) {
             expect(e.message).toBe('Cannot define a suite inside an environment');
           }
-          try {
-            t.environment('env2', () => {});
-            expect(true).toBe(false);
-          } catch (e) {
-            expect(e.message).toBe('Cannot define an environment inside an environment');
-          }
         });
+        try {
+          t.it('test', () => {}).environment(env).environment(env2);
+          expect(true).toBe(false);
+        } catch (e) {
+          expect(e.message).toBe('Cannot use environments "env2" and "env" that share a parent environment "suite1 env" in test "suite1 test"');
+        }
+        try {
+          t.it('test', () => {}).environment(env2).environment(env);
+          expect(true).toBe(false);
+        } catch (e) {
+          expect(e.message).toBe('Cannot use environments "env" and "env2" that share a parent environment "suite1 env" in test "suite1 test"');
+        }
       });
       try {
         t.it('test', () => {}).environment(env);