From cadc419928c3ae3796ec89b434cda756d5de44ba Mon Sep 17 00:00:00 2001
From: Monica Kozbial <kozbial@google.com>
Date: Wed, 5 Aug 2020 20:24:13 -0700
Subject: [PATCH] Adding cleanup to theme tests (#4112)

---
 tests/mocha/test_helpers.js | 10 ++++---
 tests/mocha/theme_test.js   | 53 ++++++++++++++++++++-----------------
 2 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/tests/mocha/test_helpers.js b/tests/mocha/test_helpers.js
index fda3aef1d..f098fe888 100644
--- a/tests/mocha/test_helpers.js
+++ b/tests/mocha/test_helpers.js
@@ -51,7 +51,8 @@ function workspaceTeardown(workspace) {
     workspace.dispose();
     this.clock.runAll();  // Run all remaining queued setTimeout calls.
   } catch (e) {
-    console.error(this.currentTest.fullTitle() + '\n', e);
+    var testRef = this.currentTest || this.test;
+    console.error(testRef.fullTitle() + '\n', e);
   }
 }
 
@@ -108,8 +109,9 @@ function sharedTestSetup(options = {}) {
  * outermost suite using sharedTestTeardown.call(this).
  */
 function sharedTestTeardown() {
+  var testRef = this.currentTest || this.test;
   if (!this.sharedSetupCalled_) {
-    console.error('"' + this.currentTest.fullTitle() +
+    console.error('"' + testRef.fullTitle() +
         '" did not call sharedTestSetup');
   }
 
@@ -121,7 +123,7 @@ function sharedTestTeardown() {
       this.clock.runAll();  // Run all queued setTimeout calls.
     }
   } catch (e) {
-    console.error(this.currentTest.fullTitle() + '\n', e);
+    console.error(testRef.fullTitle() + '\n', e);
   } finally {
     // Clear Blockly.Event state.
     Blockly.Events.setGroup(false);
@@ -131,7 +133,7 @@ function sharedTestTeardown() {
       // (i.e. a previous test added an event to the queue on a timeout that
       // did not use a stubbed clock).
       Blockly.Events.FIRE_QUEUE_.length = 0;
-      console.warn(this.currentTest.fullTitle() +
+      console.warn(testRef.fullTitle() +
           '" needed cleanup of Blockly.Events.FIRE_QUEUE_. This may indicate ' +
           'leakage from an earlier test');
     }
diff --git a/tests/mocha/theme_test.js b/tests/mocha/theme_test.js
index 36f9fc06b..240032458 100644
--- a/tests/mocha/theme_test.js
+++ b/tests/mocha/theme_test.js
@@ -11,7 +11,11 @@
 'use strict';
 
 suite('Theme', function() {
+  setup(function() {
+    sharedTestSetup.call(this);
+  });
   teardown(function() {
+    sharedTestTeardown.call(this);
     // Clear all registered themes.
     Blockly.registry.typeMap_['theme'] = {};
   });
@@ -120,37 +124,38 @@ suite('Theme', function() {
 
   test('Set Theme', function() {
     defineThemeTestBlocks();
-    var blockStyles = createBlockStyles();
-    var workspace = new Blockly.WorkspaceSvg(new Blockly.Options({}));
-    var blockA = workspace.newBlock('stack_block');
+    try {
+      var blockStyles = createBlockStyles();
+      var workspace = new Blockly.WorkspaceSvg(new Blockly.Options({}));
+      var blockA = workspace.newBlock('stack_block');
 
-    blockA.setStyle = function() {this.styleName_ = 'styleTwo';};
-    var callCount = 1;
-    workspace.refreshToolboxSelection = function() {
-      return ++callCount;
-    };
-    blockA.styleName_ = 'styleOne';
+      blockA.setStyle = function() {this.styleName_ = 'styleTwo';};
+      var refreshToolboxSelectionStub =
+          sinon.stub(workspace, 'refreshToolboxSelection');
+      blockA.styleName_ = 'styleOne';
 
-    var stub = sinon.stub(Blockly, "getMainWorkspace").returns(workspace);
-    var hideStub = sinon.stub(Blockly, "hideChaff");
+      // Stubs are cleaned up in sharedTestTeardown
+      sinon.stub(Blockly, "getMainWorkspace").returns(workspace);
+      sinon.stub(Blockly, "hideChaff");
 
-    workspace.setTheme(blockStyles);
+      workspace.setTheme(blockStyles);
 
-    // Checks that the theme was set correctly on Blockly namespace
-    stringifyAndCompare(workspace.getTheme(), blockStyles);
+      // Checks that the theme was set correctly on Blockly namespace
+      stringifyAndCompare(workspace.getTheme(), blockStyles);
 
-    // Checks that the setTheme function was called on the block
-    chai.assert.equal(blockA.getStyleName(), 'styleTwo');
+      // Checks that the setTheme function was called on the block
+      chai.assert.equal(blockA.getStyleName(), 'styleTwo');
 
-    // check that the toolbox refreshed method was called
-    chai.assert.equal(workspace.refreshToolboxSelection(), 3);
+      // Checks that the toolbox refreshed method was called
+      sinon.assert.calledOnce(refreshToolboxSelectionStub);
 
-    chai.assert.equal(Blockly.Events.FIRE_QUEUE_.pop().element, 'theme');
-
-    undefineThemeTestBlocks();
-
-    stub.restore();
-    hideStub.restore();
+      assertLastCallEventArgEquals(
+          this.eventsFireStub, Blockly.Events.UI, workspace.id,
+          null, {element: 'theme'});
+    } finally {
+      workspaceTeardown.call(this, workspace);
+      undefineThemeTestBlocks();
+    }
   });
 
   suite('Validate block styles', function() {