Merge pull request #180 from JoshKarpel/lambdas-in-each

Rewrite test functions with `lambda` in `each` correctly

ward/rewrite.py

@@ -120,8 +120,13 @@ def rewrite_assertion(test: Test) -> Test:
     if test.fn.__closure__:
         clo_glob = test.fn.__closure__[0].cell_contents.__globals__
 
+    # Look through the new module,
+    # find the code object with the same name as the original code object,
+    # and build a new function with the injected assert functions added to the global namespace.
+    # Filtering on the code object name prevents finding other kinds of code objects,
+    # like lambdas stored directly in test function arguments.
     for const in new_mod_code_obj.co_consts:
-        if isinstance(const, types.CodeType):
+        if isinstance(const, types.CodeType) and const.co_name == code_obj.co_name:
             new_test_func = types.FunctionType(
                 const,
                 {**assert_func_namespace, **test.fn.__globals__, **clo_glob},

ward/tests/test_rewrite.py

@@ -187,3 +187,19 @@ def _():
         @test("test with indentation level of 2")
         def _():
             assert 2 + 3 == 5
+
+
+@test("rewriter finds correct function when there is a lambda in an each")
+def _():
+    @testable_test
+    def _(x=each(lambda: 5)):
+        assert x == 5
+
+    t = Test(fn=_, module_name="m")
+
+    rewritten = rewrite_assertions_in_tests([t])[0]
+
+    # https://github.com/darrenburns/ward/issues/169
+    # The assertion rewriter thought the lambda function stored in co_consts was the test function,
+    # so it was rebuilding the test function using the lambda as the test instead of the original function.
+    assert rewritten.fn.__code__.co_name != "<lambda>"