diff --git a/mythril/analysis/traceexplore.py b/mythril/analysis/traceexplore.py
index f88d6410..6c694199 100644
--- a/mythril/analysis/traceexplore.py
+++ b/mythril/analysis/traceexplore.py
@@ -81,8 +81,12 @@ def get_serializable_statespace(statespace):
             if (len(code_split) < 7)
             else "\\n".join(code_split[:6]) + "\\n(click to expand +)"
         )
-
-        color = color_map[node.get_cfg_dict()["contract_name"]]
+        try:
+            color = color_map[node.get_cfg_dict()["contract_name"]]
+        except KeyError:
+            color = colors[i]
+            i += 1
+            color_map[node.get_cfg_dict()["contract_name"]] = color
 
         def get_state_accounts(node_state):
             """
@@ -97,7 +101,7 @@ def get_serializable_statespace(statespace):
                 account["balance"] = str(account["balance"])
 
                 storage = {}
-                for storage_key in account["storage"].keys():
+                for storage_key in account["storage"].printable_storage:
                     storage[str(storage_key)] = str(account["storage"][storage_key])
 
                 state_accounts.append({"address": key, "storage": storage})
diff --git a/tests/graph_test.py b/tests/graph_test.py
index 1e1e3350..ac33cabd 100644
--- a/tests/graph_test.py
+++ b/tests/graph_test.py
@@ -1,45 +1,25 @@
-from mythril.analysis.callgraph import generate_graph
+"""
+This test only checks whether dumping is successful, not whether the dumped state space makes sense
+"""
 from mythril.mythril import MythrilAnalyzer, MythrilDisassembler
 from mythril.ethereum import util
 from mythril.solidity.soliditycontract import EVMContract
-from tests import (
-    BaseTestCase,
-    TESTDATA_INPUTS,
-    TESTDATA_OUTPUTS_EXPECTED,
-    TESTDATA_OUTPUTS_CURRENT,
-)
-import re
+from tests import TESTDATA_INPUTS
 
 
-class GraphTest(BaseTestCase):
-    def test_generate_graph(self):
-        for input_file in TESTDATA_INPUTS.iterdir():
-            output_expected = TESTDATA_OUTPUTS_EXPECTED / (
-                input_file.name + ".graph.html"
-            )
-            output_current = TESTDATA_OUTPUTS_CURRENT / (
-                input_file.name + ".graph.html"
-            )
+def test_generate_graph():
+    for input_file in TESTDATA_INPUTS.iterdir():
+        if input_file.name != "origin.sol.o":
+            continue
+        contract = EVMContract(input_file.read_text())
+        disassembler = MythrilDisassembler()
+        disassembler.contracts.append(contract)
+        analyzer = MythrilAnalyzer(
+            disassembler=disassembler,
+            strategy="dfs",
+            execution_timeout=5,
+            max_depth=30,
+            address=(util.get_indexed_address(0)),
+        )
 
-            contract = EVMContract(input_file.read_text())
-            disassembler = MythrilDisassembler()
-            disassembler.contracts.append(contract)
-            analyzer = MythrilAnalyzer(
-                disassembler=disassembler,
-                strategy="dfs",
-                execution_timeout=5,
-                max_depth=30,
-                address=(util.get_indexed_address(0)),
-            )
-
-            html = analyzer.graph_html(transaction_count=1)
-            output_current.write_text(html)
-
-            lines_expected = re.findall(
-                r"'label': '.*'", str(output_current.read_text())
-            )
-            lines_found = re.findall(r"'label': '.*'", str(output_current.read_text()))
-            if not (lines_expected == lines_found):
-                self.found_changed_files(input_file, output_expected, output_current)
-
-        self.assert_and_show_changed_files()
+        analyzer.graph_html(transaction_count=1)
diff --git a/tests/statespace_test.py b/tests/statespace_test.py
new file mode 100644
index 00000000..6682ab7c
--- /dev/null
+++ b/tests/statespace_test.py
@@ -0,0 +1,23 @@
+from mythril.mythril import MythrilAnalyzer, MythrilDisassembler
+from mythril.ethereum import util
+from mythril.solidity.soliditycontract import EVMContract
+from tests import TESTDATA_INPUTS
+
+
+def test_statespace_dump():
+    for input_file in TESTDATA_INPUTS.iterdir():
+        if input_file.name not in ("origin.sol.o", "suicide.sol.o"):
+            # It's too slow, so it's better to skip some tests.
+            continue
+        contract = EVMContract(input_file.read_text())
+        disassembler = MythrilDisassembler()
+        disassembler.contracts.append(contract)
+        analyzer = MythrilAnalyzer(
+            disassembler=disassembler,
+            strategy="dfs",
+            execution_timeout=5,
+            max_depth=30,
+            address=(util.get_indexed_address(0)),
+        )
+
+        analyzer.dump_statespace(contract=contract)