codeflash-internal/codeflash/verification/verification_utils.py
2023-11-16 15:18:05 -08:00

26 lines
975 B
Python

import ast
import os
def get_test_file_path(test_dir: str, function_name: str, iteration: int = 0, test_type: str = "unit") -> str:
assert test_type in ["unit", "inspired", "replay"]
function_name = function_name.replace(".", "_")
path = os.path.join(test_dir, f"test_{function_name}__{test_type}_test_{iteration}.py")
if os.path.exists(path):
return get_test_file_path(test_dir, function_name, iteration + 1)
return path
def delete_multiple_if_name_main(test_ast: ast.Module) -> ast.Module:
if_indexes = []
for index, node in enumerate(test_ast.body):
if isinstance(node, ast.If):
if (
node.test.comparators[0].value == "__main__"
and node.test.left.id == "__name__"
and isinstance(node.test.ops[0], ast.Eq)
):
if_indexes.append(index)
for index in list(reversed(if_indexes))[1:]:
del test_ast.body[index]
return test_ast