from testgen.instrumentation.edit_generated_test import parse_module_to_cst
from testgen.postprocessing.topdef_terminator import delete_top_def_nodes


def test_delete_top_def_nodes() -> None:
    module_str = """def helper():
    pass

class MyClass:
    bob = your_uncle

def non_helper():
    pass

class ClassTypeHelper:
    def helper_method(self):
        pass
    def non_helper_method(self):
        pass

class ClassTypeNonHelper:
    def other_method(self):
        pass
"""

    deletable_list = ["module.helper", "module.ClassTypeHelper.helper_method", "module.MyClass"]
    expected = """def non_helper():
    pass

class ClassTypeNonHelper:
    def other_method(self):
        pass
"""
    result = delete_top_def_nodes(parse_module_to_cst(module_str), deletable_list).code
    assert result == expected
