From 1b02db8de10f32fba0148727637dce9319757e69 Mon Sep 17 00:00:00 2001
From: Cadu <cadu.coelho@gmail.com>
Date: Sat, 16 Nov 2024 13:49:54 -0300
Subject: [PATCH] Fix tests #2

---
 tests/test_basic.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/test_basic.py b/tests/test_basic.py
index 78ba39c..a728b40 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -40,12 +40,16 @@ def test_assert_function_actually_gets_called(mock_exit):
     # The "noqa" flag here is important, or else our pre-commit hooks (flake) will remove this assignment.
     __name__ = "__main__"  # noqa
 
-    @main
     def my_main_func():
         """
         The answer to life, the universe, and everything.
         """
         builtins.exit(42)
 
+    my_main_func.__module__ = "__main__"
+
+    # Decorate it
+    main(my_main_func)
+
     # Ensure that our main function was able to call mock_exit with the expected value.
     assert EXIT_CODE_RECEIVED == 42