From 1f2803bbc641a07a6e519be1080945f6987a70fc Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Fri, 3 Apr 2020 09:48:01 -0700
Subject: [PATCH] feat(testrunner): removeEnvironment (#1650)

---
 utils/testrunner/TestRunner.js           | 10 ++++++-
 utils/testrunner/test/testrunner.spec.js | 35 ++++++++++++++++++++----
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/utils/testrunner/TestRunner.js b/utils/testrunner/TestRunner.js
index fe4eaaee9b269..7ce45d45e834d 100644
--- a/utils/testrunner/TestRunner.js
+++ b/utils/testrunner/TestRunner.js
@@ -155,7 +155,7 @@ class Test {
     return this._hooks.filter(hook => !name || hook.name === name);
   }
 
-  environment(environment) {
+  addEnvironment(environment) {
     const parents = new Set();
     for (let parent = environment; !(parent instanceof Suite); parent = parent.parentEnvironment())
       parents.add(parent);
@@ -174,6 +174,14 @@ class Test {
     }
     throw new Error(`Cannot use environment "${environment.name()}" from suite "${environment.parentSuite().fullName()}" in unrelated test "${this.fullName()}"`);
   }
+
+  removeEnvironment(environment) {
+    const index = this._environments.indexOf(environment);
+    if (index === -1)
+      throw new Error(`Environment "${environment.name()}" cannot be removed because it was not added to the test "${test.fullName()}"`);
+    this._environments.splice(index, 1);
+    return this;
+  }
 }
 
 class Environment {
diff --git a/utils/testrunner/test/testrunner.spec.js b/utils/testrunner/test/testrunner.spec.js
index f07acddb641d0..413ff56f10825 100644
--- a/utils/testrunner/test/testrunner.spec.js
+++ b/utils/testrunner/test/testrunner.spec.js
@@ -275,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(e2);
+      t.it('cuatro', () => log.push('test #4')).addEnvironment(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(e2);
+            t.it('cinco', () => log.push('test #5')).addEnvironment(e2);
           });
         });
       });
@@ -351,6 +351,31 @@ module.exports.addTests = function({testRunner, expect}) {
         'root:afterAll2',
       ]);
     });
+    it('should remove environment', async() => {
+      const log = [];
+      const t = newTestRunner();
+      const e = 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'));
+      });
+      const e2 = t.environment('env2', () => {
+        t.beforeAll(() => log.push('env2:beforeAll'));
+        t.afterAll(() => log.push('env2:afterAll'));
+        t.beforeEach(() => log.push('env2:beforeEach'));
+        t.afterEach(() => log.push('env2:afterEach'));
+      });
+      t.it('uno', () => log.push('test #1')).addEnvironment(e).addEnvironment(e2).removeEnvironment(e);
+      await t.run();
+      expect(log).toEqual([
+        'env2:beforeAll',
+        'env2:beforeEach',
+        'test #1',
+        'env2:afterEach',
+        'env2:afterAll',
+      ]);
+    });
     it('environment restrictions', async () => {
       const t = newTestRunner();
       let env;
@@ -372,20 +397,20 @@ module.exports.addTests = function({testRunner, expect}) {
           }
         });
         try {
-          t.it('test', () => {}).environment(env).environment(env2);
+          t.it('test', () => {}).addEnvironment(env).addEnvironment(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);
+          t.it('test', () => {}).addEnvironment(env2).addEnvironment(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);
+        t.it('test', () => {}).addEnvironment(env);
         expect(true).toBe(false);
       } catch (e) {
         expect(e.message).toBe('Cannot use environment "env" from suite "suite1" in unrelated test "test"');