create virtual environment installs the backend s package and common requirements ensure other backends are uninstalled install whl package runs the example script exits virtual environment deletes files and any miscellaneous install logs ensures packages from all backends are installed builds keras core package and returns package file path creates and activates a virtual environment ensures the backend s package is installed and the other backends are uninstalled runs test of basic flow in keras core tests for backendspecific imports and model fit removes virtual environment and associated files create virtual environment installs the backend s package and common requirements ensure other backends are uninstalled install whl package runs the example script exits virtual environment deletes files and any miscellaneous install logs ensures packages from all backends are installed builds keras core package and returns package file path creates and activates a virtual environment ensures the backend s package is installed and the other backends are uninstalled runs test of basic flow in keras core tests for backend specific imports and model fit removes virtual environment and associated files
import os
import re
import subprocess

from keras import backend

BACKEND_REQ = {
    "tensorflow": "tensorflow",
    "torch": "torch torchvision",
    "jax": "jax jaxlib",
}


def setup_package():
    subprocess.run("rm -rf tmp_build_dir", shell=True)
    build_process = subprocess.run(
        "python3 pip_build.py",
        capture_output=True,
        text=True,
        shell=True,
    )
    print(build_process.stdout)
    whl_path = re.findall(
        r"[^\s]*\.whl",
        build_process.stdout,
    )[-1]
    if not whl_path:
        print(build_process.stderr)
        raise ValueError("Installing Keras package unsuccessful. ")
    return whl_path


def create_virtualenv():
    env_setup = [
        
        "python3 -m venv test_env",
    ]
    os.environ["PATH"] = (
        "/test_env/bin/" + os.pathsep + os.environ.get("PATH", "")
    )
    run_commands_local(env_setup)


def manage_venv_installs(whl_path):
    other_backends = list(set(BACKEND_REQ.keys()) - {backend.backend()})
    install_setup = [
        
        "pip install " + BACKEND_REQ[backend.backend()],
        "pip install -r requirements-common.txt",
        "pip install pytest",
        
        "pip uninstall -y "
        + BACKEND_REQ[other_backends[0]]
        + " "
        + BACKEND_REQ[other_backends[1]],
        
        "pip install " + whl_path,
    ]
    run_commands_venv(install_setup)


def run_keras_flow():
    test_script = [
        
        "python -m pytest integration_tests/basic_full_flow.py",
    ]
    run_commands_venv(test_script)


def cleanup():
    cleanup_script = [
        
        
        "exit",
        "rm -rf test_env",
        "rm -rf tmp_build_dir",
        "rm -f *+cpu",
    ]
    run_commands_local(cleanup_script)


def run_commands_local(commands):
    for command in commands:
        print(f"Running command: {command}")
        subprocess.run(command, shell=True)


def run_commands_venv(commands):
    for command in commands:
        print(f"Running command: {command}")
        cmd_with_args = command.split(" ")
        cmd_with_args[0] = "test_env/bin/" + cmd_with_args[0]
        p = subprocess.Popen(cmd_with_args)
        assert p.wait() == 0


def test_keras_imports():
    try:
        
        
        whl_path = setup_package()

        
        create_virtualenv()

        
        
        manage_venv_installs(whl_path)

        
        
        run_keras_flow()

        
    finally:
        cleanup()


if __name__ == "__main__":
    test_keras_imports()