test-unit

#!/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 shutil
import subprocess
import sys

from shared import utils

PYTHON_TEST_SUITES = [
    {
        "name": "sandbox-router",
        "dir": os.path.join("clients", "python", "agentic-sandbox-client", "sandbox-router"),
        "requirements": "requirements.txt",
    },
    {
        "name": "k8s-agent-sandbox",
        "dir": os.path.join("clients", "python", "agentic-sandbox-client", "k8s_agent_sandbox"),
        "requirements": "requirements.txt",
    },
    {
        "name": "dev-tools-shared",
        "dir": os.path.join("dev", "tools", "shared"),
        # No requirements file; headers.py relies only on the standard library,
        # so the runner just installs pytest and runs the suite.
        "requirements": "requirements.txt",
    },
    {
        "name": "agent-sandbox-rl",
        "dir": os.path.join("examples", "agent-sandbox-rl"),
        # Editable installs (in order): the in-repo SDK first (so we don't depend
        # on a PyPI build), then the example package with its test extra. The
        # mocked suite needs neither datasets nor r2egym (both stubbed in tests).
        "editable_installs": [
            (os.path.join("clients", "python", "agentic-sandbox-client"), ""),
            (os.path.join("examples", "agent-sandbox-rl"), "[test]"),
        ],
    },
]


def run_go_tests(repo_root):
    """ runs Go unit tests and returns the exit code """
    go_list_cmd = ["go", "list", "./..."]
    go_list_output = subprocess.check_output(go_list_cmd, cwd=repo_root, text=True)
    packages = go_list_output.strip().split('\n')

    filtered_packages = [pkg for pkg in packages if "test/e2e" not in pkg]

    result = subprocess.run(utils.go_tool_args(
        "gotestsum",
        f"--junitfile={repo_root}/bin/unit-junit.xml",
        "--",
        "-race",
        *filtered_packages
    ), cwd=repo_root)

    return result.returncode


def run_python_tests(repo_root):
    """ runs Python unit tests in a temporary venv and returns the exit code """
    returncode = 0

    for suite in PYTHON_TEST_SUITES:
        test_dir = os.path.join(repo_root, suite["dir"])
        if not os.path.isdir(test_dir):
            print(f"WARNING: Python test directory not found, skipping: {test_dir}")
            continue

        venv_dir = os.path.join(repo_root, "bin", f"python-venv-{suite['name']}")
        venv_python = os.path.join(venv_dir, "bin", "python")
        venv_pip = os.path.join(venv_dir, "bin", "pip")

        print(f"\n=== Setting up Python venv for {suite['name']} tests ===")
        if os.path.isdir(venv_dir):
            # Safely verify the directory is indeed a virtual environment before removing
            if os.path.exists(os.path.join(venv_dir, "pyvenv.cfg")):
                shutil.rmtree(venv_dir)
            else:
                print(f"WARNING: Directory {venv_dir} exists but does not appear to be a venv (missing pyvenv.cfg). Skipping clean.")
        subprocess.check_call([sys.executable, "-m", "venv", venv_dir])

        env = os.environ.copy()
        if not os.environ.get("PIP_EXTRA_INDEX_URL"):
            env.pop("PIP_EXTRA_INDEX_URL", None)

        req_name = suite.get("requirements")
        req_file = os.path.join(test_dir, req_name) if req_name else None
        if suite.get("editable_installs"):
            # install each (relative_dir, extras) editable, in order
            for rel_dir, extras in suite["editable_installs"]:
                subprocess.check_call(
                    [venv_pip, "install", "--quiet", "-e", f".{extras}"],
                    cwd=os.path.join(repo_root, rel_dir), env=env)
        elif req_file and os.path.isfile(req_file):
            subprocess.check_call(
                [venv_pip, "install", "--quiet", "-r", req_file], env=env)
        elif suite["name"] == "k8s-agent-sandbox":
            # install the package itself with test extras
            package_dir = os.path.dirname(test_dir)
            subprocess.check_call(
            [venv_pip, "install", "--quiet", "-e", ".[test]"], cwd=package_dir, env=env)

        subprocess.check_call(
            [venv_pip, "install", "--quiet", "pytest"], env=env)

        print(f"\n=== Running Python unit tests: {suite['name']} ===")
        junit_file = os.path.join(
            repo_root, "bin", f"python-{suite['name']}-junit.xml")
        result = subprocess.run([
            venv_python, "-m", "pytest",
            test_dir,
            f"--junitxml={junit_file}",
            "-v",
        ], cwd=test_dir)

        if result.returncode != 0:
            returncode = result.returncode

    return returncode


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

    go_rc = run_go_tests(repo_root)
    python_rc = run_python_tests(repo_root)

    if go_rc != 0:
        return go_rc
    return python_rc


if __name__ == "__main__":
    sys.exit(main())