diff --git a/src/App.tsx b/src/App.tsx
index 8671cc74..41b36686 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -134,7 +134,7 @@ const App: React.FC = () => {
 
     openStorage();
     initializeBlocks();
-    //testAllBlocksInToolbox(toolbox.getToolboxJSON([], []).contents);
+    //testAllBlocksInToolbox();
   }, []);
 
   const openStorage = async () => {
diff --git a/src/toolbox/toolbox.ts b/src/toolbox/toolbox.ts
index 9f261611..ce8c361a 100644
--- a/src/toolbox/toolbox.ts
+++ b/src/toolbox/toolbox.ts
@@ -31,7 +31,7 @@ import {category as methodsCategory} from './methods_category';
 
 export function getToolboxJSON(
     opt_includeExportedBlocksFromProject: toolboxItems.ContentsType[],
-    shownPythonToolboxCategories: Set<string>) {
+    shownPythonToolboxCategories: Set<string> | null) {
   const contents: toolboxItems.ContentsType[] = generatedToolbox.getToolboxCategories();
   filterGeneratedCategories(contents, shownPythonToolboxCategories);
 
@@ -81,7 +81,7 @@ export function getToolboxJSON(
 }
 
 function filterGeneratedCategories(
-    contents: toolboxItems.ContentsType[], shownPythonToolboxCategories: Set<string>) {
+    contents: toolboxItems.ContentsType[], shownPythonToolboxCategories: Set<string> | null) {
   contents.forEach((item) => {
     if (item.kind === 'category') {
       const category = item as toolboxItems.Category;
@@ -91,7 +91,7 @@ function filterGeneratedCategories(
       }
       if ((category as toolboxItems.PythonModuleCategory).moduleName) {
         const moduleName = (item as toolboxItems.PythonModuleCategory).moduleName;
-        if (!shownPythonToolboxCategories.has(moduleName)) {
+        if (shownPythonToolboxCategories != null && !shownPythonToolboxCategories.has(moduleName)) {
           if (category.contents) {
             removeBlocksAndSeparators(category.contents);
           }
@@ -99,7 +99,7 @@ function filterGeneratedCategories(
       }
       if ((category as toolboxItems.PythonClassCategory).className) {
         const className = (item as toolboxItems.PythonClassCategory).className;
-        if (!shownPythonToolboxCategories.has(className)) {
+        if (shownPythonToolboxCategories != null && !shownPythonToolboxCategories.has(className)) {
           if (category.contents) {
             removeBlocksAndSeparators(category.contents);
           }
@@ -123,7 +123,7 @@ function removeBlocksAndSeparators(contents: toolboxItems.ContentsType[]) {
 }
 
 function removeEmptyCategories(
-    contents: toolboxItems.ContentsType[], shownPythonToolboxCategories: Set<string>) {
+    contents: toolboxItems.ContentsType[], shownPythonToolboxCategories: Set<string> | null) {
   let i = 0;
   while (i < contents.length) {
     let remove = false;
@@ -137,7 +137,7 @@ function removeEmptyCategories(
       }
       if (category.contents &&
           category.contents.length == 0 &&
-          !shownPythonToolboxCategories.has(fullCategoryName)) {
+          shownPythonToolboxCategories != null && !shownPythonToolboxCategories.has(fullCategoryName)) {
         remove = true;
       }
     }
diff --git a/src/toolbox/toolbox_tests.ts b/src/toolbox/toolbox_tests.ts
index 5aa00f0d..b67e953b 100644
--- a/src/toolbox/toolbox_tests.ts
+++ b/src/toolbox/toolbox_tests.ts
@@ -20,13 +20,15 @@
  */
 
 import * as Blockly from 'blockly/core';
-import { pythonGenerator } from 'blockly/python';
+import { extendedPythonGenerator } from '../editor/extended_python_generator';
 import * as toolboxItems from './items';
+import { getToolboxJSON } from '../toolbox/toolbox';
 
 
 // Tests
 
-export function testAllBlocksInToolbox(contents: toolboxItems.ContentsType[]) {
+export function testAllBlocksInToolbox() {
+  const contents: toolboxItems.ContentsType[] = getToolboxJSON([], null).contents;
   alert('Press OK to run tests on all blocks from the toolbox.');
   const toolboxTestData = new ToolboxTestData(contents, () => {
     alert('Completed tests on all blocks in the toolbox. See console for any errors.');
@@ -104,24 +106,24 @@ class ToolboxTestData {
   }
 
   private testBlockPython(block: Blockly.Block) {
-    pythonGenerator.init(this.blocklyWorkspace);
+    extendedPythonGenerator.init(this.blocklyWorkspace);
     let code = null;
     try {
-      code = pythonGenerator.blockToCode(block);
+      code = extendedPythonGenerator.blockToCode(block);
     } catch (e) {
       console.log('Error - ' + e + ' - Unable to test block...'); console.log(block);
     }
     if (block.outputConnection) {
       if (!Array.isArray(code)) {
         console.log(
-            'Error - pythonGenerator.forBlock["' + block.type + '"] ' +
+            'Error - extendedPythonGenerator.forBlock["' + block.type + '"] ' +
             'generated a ' + (typeof code) +
             ', but it should generate an array because the block has an output connection.');
       }
     } else {
       if (typeof code != 'string') {
         console.log(
-            'Error - pythonGenerator.forBlock["' + block.type + '"] ' +
+            'Error - extendedPythonGenerator.forBlock["' + block.type + '"] ' +
             'generated a ' + (typeof code) +
             ', but should generate a string because the block does not have an output connection.');
       }