From d15ff852ee3f9dd98fb74766ee2a9b522dffef33 Mon Sep 17 00:00:00 2001
From: Nikhil Parasaram <tommycjniko@gmail.com>
Date: Sat, 2 Feb 2019 21:27:42 +0530
Subject: [PATCH] Add graph analysis test

---
 mythril/mythril/mythril_analyzer.py     | 11 +++++++----
 mythril/mythril/mythril_disassembler.py |  2 +-
 tests/graph_test.py                     | 13 +++++++------
 tests/mythril/mythril_analyzer_test.py  |  2 --
 4 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/mythril/mythril/mythril_analyzer.py b/mythril/mythril/mythril_analyzer.py
index d1f9b807..2905fbe3 100644
--- a/mythril/mythril/mythril_analyzer.py
+++ b/mythril/mythril/mythril_analyzer.py
@@ -49,7 +49,7 @@ class MythrilAnalyzer:
     def dump_statespace(
         self,
         strategy: str,
-        contract: List[EVMContract],
+        contract: List[EVMContract] = None,
         address: Optional[str] = None,
         max_depth: Optional[int] = None,
         execution_timeout: Optional[int] = None,
@@ -68,7 +68,7 @@ class MythrilAnalyzer:
         :return: The serialized state space
         """
         sym = SymExecWrapper(
-            contract,
+            contract or self.contracts[0],
             address,
             strategy,
             dynloader=DynLoader(
@@ -87,13 +87,14 @@ class MythrilAnalyzer:
     def graph_html(
         self,
         strategy: str,
-        contract: List[EVMContract],
         address: str,
+        contract: List[EVMContract] = None,
         max_depth: Optional[int] = None,
         enable_physics: bool = False,
         phrackify: bool = False,
         execution_timeout: Optional[int] = None,
         create_timeout: Optional[int] = None,
+        transaction_count: Optional[int] = None,
         enable_iprof: bool = False,
     ) -> str:
         """
@@ -106,11 +107,12 @@ class MythrilAnalyzer:
         :param phrackify: If true generates Phrack-style call graph
         :param execution_timeout: The total execution timeout of the contract
         :param create_timeout: The total contract creation timeout
+        :param transaction_count: The amount of transactions to be executed
         :param enable_iprof: Enables/disables instruction profiler
         :return: The generated graph in html format
         """
         sym = SymExecWrapper(
-            contract,
+            contract or self.contracts[0],
             address,
             strategy,
             dynloader=DynLoader(
@@ -120,6 +122,7 @@ class MythrilAnalyzer:
             ),
             max_depth=max_depth,
             execution_timeout=execution_timeout,
+            transaction_count=transaction_count,
             create_timeout=create_timeout,
             enable_iprof=enable_iprof,
         )
diff --git a/mythril/mythril/mythril_disassembler.py b/mythril/mythril/mythril_disassembler.py
index 3e455115..a55acf5d 100644
--- a/mythril/mythril/mythril_disassembler.py
+++ b/mythril/mythril/mythril_disassembler.py
@@ -21,7 +21,7 @@ log = logging.getLogger(__name__)
 class MythrilDisassembler:
     def __init__(
         self,
-        eth: Optional[EthJsonRpc],
+        eth: Optional[EthJsonRpc] = None,
         solc_version: str = None,
         solc_args: str = None,
         enable_online_lookup: bool = False,
diff --git a/tests/graph_test.py b/tests/graph_test.py
index ba84f7da..35a7fd32 100644
--- a/tests/graph_test.py
+++ b/tests/graph_test.py
@@ -1,5 +1,5 @@
 from mythril.analysis.callgraph import generate_graph
-from mythril.analysis.symbolic import SymExecWrapper
+from mythril.mythril import MythrilAnalyzer, MythrilDisassembler
 from mythril.ethereum import util
 from mythril.solidity.soliditycontract import EVMContract
 from tests import (
@@ -22,16 +22,17 @@ class GraphTest(BaseTestCase):
             )
 
             contract = EVMContract(input_file.read_text())
+            disassembler = MythrilDisassembler()
+            disassembler.contracts.append(contract)
+            analyzer = MythrilAnalyzer(disassembler)
 
-            sym = SymExecWrapper(
-                contract,
-                address=(util.get_indexed_address(0)),
+            html = analyzer.graph_html(
                 strategy="dfs",
                 transaction_count=1,
                 execution_timeout=5,
+                max_depth=30,
+                address=(util.get_indexed_address(0)),
             )
-
-            html = generate_graph(sym)
             output_current.write_text(html)
 
             lines_expected = re.findall(
diff --git a/tests/mythril/mythril_analyzer_test.py b/tests/mythril/mythril_analyzer_test.py
index 90423dcc..a0220893 100644
--- a/tests/mythril/mythril_analyzer_test.py
+++ b/tests/mythril/mythril_analyzer_test.py
@@ -4,8 +4,6 @@ from mythril.mythril import MythrilDisassembler, MythrilAnalyzer
 
 def test_fire_lasers():
     disassembler = MythrilDisassembler(eth=None)
-    print(str((Path(__file__).parent.parent / "testdata/input_contracts/origin.sol")))
-    print(Path(__file__))
     disassembler.load_from_solidity(
         [
             str(