From b26efc797d70b7e3fb0986faa941acbf6511ad80 Mon Sep 17 00:00:00 2001
From: "Johan B.W. de Vries" <info@jbwdevries.nl>
Date: Sat, 7 Aug 2021 14:52:51 +0200
Subject: [PATCH] Cleanup, added tests

---
 py2wasm/python.py                | 31 +++++++++++++------------------
 tests/integration/test_simple.py | 18 +++++++++++++++++-
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/py2wasm/python.py b/py2wasm/python.py
index 90f97af..9c6e989 100644
--- a/py2wasm/python.py
+++ b/py2wasm/python.py
@@ -295,28 +295,23 @@ 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)
+
+        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
 
diff --git a/tests/integration/test_simple.py b/tests/integration/test_simple.py
index 17b2b86..af73c4d 100644
--- a/tests/integration/test_simple.py
+++ b/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