test-e2e

#!/usr/bin/env python3
# Copyright 2025 The Kubernetes Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
import subprocess
import sys

from shared import utils


def run_go_e2e_tests(repo_root, artifact_dir, env):
    """Runs the Go e2e tests"""
    print("========= Running Go E2E tests... =========")
    race_enabled = env.get("RACE", "false").lower() in ("true", "1", "yes")
    print(f"Race detector enabled: {race_enabled}")

    go_test_args = [
       "./test/e2e/...",
       "--parallel=1",
       "-v",
    ]

    if race_enabled:
        go_test_args.append("-race")

    go_test_cmd = utils.go_tool_args(
        "gotestsum",
         f"--junitfile={repo_root}/bin/e2e-go-junit.xml",
         f"--jsonfile={artifact_dir}/test-e2e.json",
         "--",
         *go_test_args,
)
    print(f"Running command: {' '.join(go_test_cmd)}")
    return subprocess.run(go_test_cmd, env=env, cwd=repo_root).returncode


def run_go_e2e_benchmarks(repo_root, artifact_dir, env):
    """Runs the Go e2e benchmarks and saves output for benchstat"""
    print("========= Running Go E2E benchmarks... =========")
    benchmark_file = os.path.join(artifact_dir, "benchmarks-go.txt")

    # Run benchmarks with go test directly (not gotestsum) to get clean benchmark output
    go_test_cmd = [
        "go", "test",
        "-bench=.",
        "-benchtime=1x",  # Run each benchmark once (e2e tests are expensive)
        "-run=^$",  # Don't run tests, only benchmarks
        "./test/e2e/...",
    ]
    print(f"Running command: {' '.join(go_test_cmd)}")
    print(f"Benchmark output will be saved to: {benchmark_file}")

    with open(benchmark_file, "w") as f:
        process = subprocess.Popen(go_test_cmd, env=env, cwd=repo_root, stdout=subprocess.PIPE, text=True)

        for line in process.stdout:
            sys.stdout.write(line)  # Print to terminal as it happens
            f.write(line)           # Save to file also

        process.wait()

    if process.returncode != 0:
        print(f"Go E2E benchmarks process failed with exit code {process.returncode}")
        return process.returncode

    csv_file = os.path.join(artifact_dir, "benchmarks.csv")
    benchstat_cmd = utils.go_tool_args(
        "benchstat",
        "-format", "csv",
        benchmark_file
    )
    print(f"Running benchstat: {' '.join(benchstat_cmd)}")
    with open(csv_file, "w") as f_csv:
        benchstat_result = subprocess.run(benchstat_cmd, env=env, cwd=repo_root, stdout=f_csv)

    if benchstat_result.returncode != 0:
        print(f"benchstat failed with exit code {benchstat_result.returncode}")
        return benchstat_result.returncode

    return 0


def setup_python_sdk(repo_root, env):
    """Installs the Python SDK and its dependencies"""
    print("========= Ensuring Python SDK is installed... =========")
    venv_path = os.path.join(repo_root, ".venv")
    sdk_path = os.path.join(repo_root, "clients", "python", "agentic-sandbox-client")

    # Create virtual environment
    print(f"Creating virtual environment at {venv_path}")
    result = subprocess.run(
        [sys.executable, "-m", "venv", venv_path], env=env, cwd=repo_root, check=False
    )
    if result.returncode != 0:
        print("Failed to create virtual environment.")
        return False

    # Install in editable mode within the virtual environment
    pip_executable = os.path.join(venv_path, "bin", "pip")
    install_cmd = [pip_executable, "install", "-e", f"{sdk_path}[test]"]
    # Install in editable mode
    print(f"Running command: {' '.join(install_cmd)}")
    result = subprocess.run(install_cmd, env=env, cwd=repo_root, check=False)
    if result.returncode != 0:
        print("Failed to install Python SDK.")
        return False
    print("Python SDK installed successfully.")
    return True


def run_python_e2e_tests(repo_root, artifact_dir, env):
    """Runs the Python e2e tests and generates a junit report"""
    print("========= Running Python SDK E2E tests... =========")

    # Define the path for the JUnit report
    junit_report_path = f"{repo_root}/bin/e2e-python-sdk-junit.xml"

    # Set the environment variable for the test namespace
    test_env = env.copy()

    python_executable = os.path.join(repo_root, ".venv", "bin", "python")
    pytest_cmd = [
        python_executable,
        "-m",
        "pytest",
        f"--junitxml={junit_report_path}",
        "-v",
        "-n", "auto",  # to run tests in parallel
        os.path.join(repo_root, "test", "e2e/"),
    ]

    print(f"Running command: {' '.join(pytest_cmd)}")
    result = subprocess.run(pytest_cmd, env=test_env, cwd=repo_root)
    return result.returncode


def main(args):
    """ invokes unit tests and outputs a junit results file """
    repo_root = utils.get_repo_root()

    artifact_dir = os.getenv("ARTIFACTS")
    if not artifact_dir:
        artifact_dir = f"{repo_root}/bin"
    os.makedirs(artifact_dir, exist_ok=True)

    env = os.environ.copy()
    env["IMAGE_TAG"] = utils.get_image_tag()
    env["IMAGE_PREFIX"] = utils.get_image_prefix(args)
    if not os.environ.get("PIP_EXTRA_INDEX_URL"):
        env.pop("PIP_EXTRA_INDEX_URL", None)

    if not setup_python_sdk(repo_root, env):
        return 1

    exit_code = 0

    go_test_result = 0
    go_benchmark_result = 0
    python_test_result = 0

    if args.suite in ["all", "tests"]:
        go_test_result = run_go_e2e_tests(repo_root, artifact_dir, env)
        python_test_result = run_python_e2e_tests(repo_root, artifact_dir, env)
        
    if args.suite in ["all", "benchmarks"]:
        go_benchmark_result = run_go_e2e_benchmarks(repo_root, artifact_dir, env)

    if go_test_result != 0:
        print("Go E2E tests failed.")

    if go_benchmark_result != 0:
        print("Go E2E benchmarks failed.")

    if python_test_result != 0:
        print("Python SDK E2E tests failed.")
        exit_code = 1

    if go_test_result != 0 or python_test_result != 0 or go_benchmark_result != 0:
        print("One or more E2E tests failed.")
        exit_code = 1

    return exit_code


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Run end-to-end tests against Kubernetes kind cluster"
    )
    parser.add_argument(
        "--image-prefix",
        dest="image_prefix",
        help="prefix for the image name. requires slash at the end if a path",
        type=str,
        default=os.getenv("IMAGE_PREFIX"),
    )
    parser.add_argument(
        "--suite",
        dest="suite",
        help="which suite of tests to run",
        type=str,
        choices=["all", "tests", "benchmarks"],
        default="all",
    )
    args = parser.parse_args()
    sys.exit(main(args))