From b84319b333c82038b844a8417149c75f8bb71d72 Mon Sep 17 00:00:00 2001 From: Dillon Laird Date: Mon, 14 Oct 2024 18:00:24 -0700 Subject: [PATCH] added if guard test case --- tests/unit/test_vac.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/unit/test_vac.py b/tests/unit/test_vac.py index 4f8ead55..d5cebbb5 100644 --- a/tests/unit/test_vac.py +++ b/tests/unit/test_vac.py @@ -141,3 +141,16 @@ def check_helmets(image_path): # The function can be called with the image path""" code_out = strip_function_calls(code, exclusions=["register_heif_opener"]) assert code_out == expected_code + + +def test_strip_function_call_if_name_equal_main(): + code = """import os +def f(): + print("Hello!") +if __name__ == "__main__": + f()""" + expected_code = """import os +def f(): + print("Hello!")""" + code_out = strip_function_calls(code) + assert code_out == expected_code