diff --git a/tests/solidity_contract_test.py b/tests/solidity_contract_test.py new file mode 100644 index 00000000..07781a46 --- /dev/null +++ b/tests/solidity_contract_test.py @@ -0,0 +1,28 @@ +from unittest import TestCase +from pathlib import Path + +from mythril.ether.soliditycontract import SolidityContract + +TEST_FILES = Path.cwd() / "testdata" + +class SolidityContractTest(TestCase): + + def test_get_source_info_without_name_gets_latest_contract_info(self): + input_file = TEST_FILES / "multi_contracts.sol" + contract = SolidityContract(str(input_file)) + + code_info = contract.get_source_info(142) + + self.assertEqual(code_info.filename, str(input_file)) + self.assertEqual(code_info.lineno, 14) + self.assertEqual(code_info.code, "msg.sender.transfer(2 ether)") + + def test_get_source_info_with_contract_name_specified(self): + input_file = TEST_FILES / "multi_contracts.sol" + contract = SolidityContract(str(input_file), name="Transfer1") + + code_info = contract.get_source_info(142) + + self.assertEqual(code_info.filename, str(input_file)) + self.assertEqual(code_info.lineno, 6) + self.assertEqual(code_info.code, "msg.sender.transfer(1 ether)")