Cleanup, added tests

py2wasm/python.py

@@ -295,29 +295,24 @@ class Visitor:
 
         called_name = node.func.id
 
+        search_list: List[Union[wasm.Function, wasm.Import]]
+        search_list = [
+            *module.functions,
+            *module.imports,
+        ]
+
         called_func_list = [
             x
-            for x in module.functions
+            for x in search_list
             if x.name == called_name
         ]
-        if called_func_list:
-            assert len(called_func_list) == 1
-            called_params = called_func_list[0].params
-            called_result = called_func_list[0].result
-        else:
-            called_import_list = [
-                x
-                for x in module.imports
-                if x.name == node.func.id
-            ]
-            if called_import_list:
-                assert len(called_import_list) == 1
-                called_params = called_import_list[0].params
-                called_result = called_import_list[0].result
-            else:
+
         assert 1 == len(called_func_list), \
             'Could not find function {}'.format(node.func.id)
 
+        called_params = called_func_list[0].params
+        called_result = called_func_list[0].result
+
         assert exp_type == called_result
 
         assert len(called_params) == len(node.args), \

tests/integration/test_simple.py

@@ -113,7 +113,7 @@ def testEntry(a: i32) -> i32:
     assert 0 == suite.run_code(-1).returned_value
 
 @pytest.mark.integration_test
-def test_call():
+def test_call_pre_defined():
     code_py = """
 def helper(left: i32, right: i32) -> i32:
     return left + right
@@ -127,3 +127,19 @@ def testEntry() -> i32:
 
     assert 13 == result.returned_value
     assert [] == result.log_int32_list
+
+@pytest.mark.integration_test
+def test_call_post_defined():
+    code_py = """
+@exported
+def testEntry() -> i32:
+    return helper(10, 3)
+
+def helper(left: i32, right: i32) -> i32:
+    return left - right
+"""
+
+    result = Suite(code_py, 'test_call').run_code()
+
+    assert 7 == result.returned_value
+    assert [] == result.log_int32_list