From 247cdc4c33a1dabc2b290e413cf77cd7bc9c88b6 Mon Sep 17 00:00:00 2001
From: Joran Honig <joran.honig@gmail.com>
Date: Tue, 23 Oct 2018 14:32:58 +0200
Subject: [PATCH] get_function_info tests

---
 mythril/disassembler/disassembly.py |  4 +-
 tests/disassembler/__init__.py      |  0
 tests/disassembler/disassembly.py   | 61 +++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 tests/disassembler/__init__.py
 create mode 100644 tests/disassembler/disassembly.py

diff --git a/mythril/disassembler/disassembly.py b/mythril/disassembler/disassembly.py
index 48682f5a..640c56cb 100644
--- a/mythril/disassembler/disassembly.py
+++ b/mythril/disassembler/disassembly.py
@@ -37,7 +37,7 @@ class Disassembly(object):
         )
 
         for index in jump_table_indices:
-            function_hash, jump_target, function_name = _get_function_info(
+            function_hash, jump_target, function_name = get_function_info(
                 index, self.instruction_list, signatures
             )
             self.func_hashes.append(function_hash)
@@ -52,7 +52,7 @@ class Disassembly(object):
         return asm.instruction_list_to_easm(self.instruction_list)
 
 
-def _get_function_info(index: int, instruction_list: list, signature_database: SignatureDb) -> (str, int, str):
+def get_function_info(index: int, instruction_list: list, signature_database: SignatureDb) -> (str, int, str):
     """
     Finds the function information for a call table entry
     Solidity uses the first 4 bytes of the calldata to indicate which function the message call should execute
diff --git a/tests/disassembler/__init__.py b/tests/disassembler/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/disassembler/disassembly.py b/tests/disassembler/disassembly.py
new file mode 100644
index 00000000..49bca31d
--- /dev/null
+++ b/tests/disassembler/disassembly.py
@@ -0,0 +1,61 @@
+from mythril.disassembler.disassembly import *
+
+instruction_list = [
+    {"opcode": "PUSH4", "argument": "0x10203040"},
+    {"opcode": "EQ"},
+    {"opcode": "PUSH4", "argument": "0x40302010"},
+    {"opcode": "JUMPI"},
+]
+
+
+def test_get_function_info(mocker):
+    # Arrange
+    global instruction_list
+
+    signature_database_mock = SignatureDb()
+    mocker.patch.object(signature_database_mock, "get")
+    signature_database_mock.get.return_value = ["function_name"]
+
+    # Act
+    function_hash, entry_point, function_name = get_function_info(
+        0, instruction_list, signature_database_mock
+    )
+
+    # Assert
+    assert function_hash == "0x10203040"
+    assert entry_point == 0x40302010
+    assert function_name == "function_name"
+
+
+def test_get_function_info_multiple_names(mocker):
+    # Arrange
+    global instruction_list
+
+    signature_database_mock = SignatureDb()
+    mocker.patch.object(signature_database_mock, "get")
+    signature_database_mock.get.return_value = ["function_name", "another_name"]
+
+    # Act
+    function_hash, entry_point, function_name = get_function_info(
+        0, instruction_list, signature_database_mock
+    )
+
+    # Assert
+    assert function_name == "**ambiguous** function_name"
+
+
+def test_get_function_info_no_names(mocker):
+    # Arrange
+    global instruction_list
+
+    signature_database_mock = SignatureDb()
+    mocker.patch.object(signature_database_mock, "get")
+    signature_database_mock.get.return_value = []
+
+    # Act
+    function_hash, entry_point, function_name = get_function_info(
+        0, instruction_list, signature_database_mock
+    )
+
+    # Assert
+    assert function_name == "_function_0x10203040"