Function main

codegen/test_clean_code.py:89–130

()

def main():
    code, success = clean_code(bad_code)

    if not success:
        print("Code cleaning failed!")
        return

    print(f"Cleaned code: {code}")

    # Add this call back into the code
    code += "\ndostuff()"

    # Test script file name to execute in the Docker container
    sources_dirname = "test_sources"
    test_script = "test_script.py"
    expected_output = """5 + 2 = 7
5 - 2 = 3
5 * 2 = 10
5 / 2 = 2.5
The area of a circle with radius 3 is 28.27
The square root of 49 is 7.0
The maximum value in the list is 91
7 is a prime number"""

    script_path = os.path.join(os.getcwd(), sources_dirname, test_script)
    os.makedirs(os.path.join(os.getcwd(), sources_dirname), exist_ok=True)

    # Write a simple test script to execute in the container
    with open(script_path, "w") as f:
        f.write(code)

    docker_execute = DockerExecute(sources_dirname=sources_dirname)

    # Prime the pump
    exit_code, output = docker_execute.execute(test_script)
    assert output == expected_output, f"Unexpected output: {output}"
    assert exit_code == 0, f"Unexpected exit code: {exit_code}"

    # Remove the test script file
    os.remove(script_path)

    print("Success!  Script now runs and prints the expected output.")

if __name__ == "__main__":
    main()