load("@rules_python//python:defs.bzl", "py_library")
load("//bazel:python.bzl", "py_test_module_list")

py_library(
    name = "conftest",
    srcs = glob(["**/conftest.py"]),
    visibility = [
        "//python/ray/llm/tests:__subpackages__",
    ],
)

# Common tests
py_test_module_list(
    size = "small",
    files = glob(["common/**/test_*.py"]),
    tags = [
        "cpu",
        "exclusive",
        "team:llm",
    ],
    deps = ["//:ray_lib"],
)

# Batch test
py_test_module_list(
    size = "medium",
    files = glob([
        "batch/cpu/**/test_*.py",
        "batch/observability/usage_telemetry/test_*.py",
    ]),
    tags = [
        "cpu",
        "exclusive",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)

py_test_module_list(
    size = "large",
    env = {
        "VLLM_FLASH_ATTN_VERSION": "2",
    },
    files = glob(
        ["batch/gpu/**/test_*.py"],
    ),
    tags = [
        "exclusive",
        "gpu",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)

# Serve test
# Utils unit tests
py_test_module_list(
    size = "small",
    files = glob(["serve/utils/test_*.py"]),
    tags = [
        "cpu",
        "exclusive",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)

# CPU tests
py_test_module_list(
    size = "large",
    data = glob(["serve/**/*.yaml"]),
    # *_direct_streaming.py files run only in the dedicated target below,
    # which sets the direct-streaming + HAProxy env they require.
    files = glob(
        ["serve/cpu/**/test_*.py"],
        exclude = ["serve/cpu/**/*_direct_streaming.py"],
    ),
    tags = [
        "cpu",
        "exclusive",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)

# Direct-streaming session-affinity tests
py_test_module_list(
    size = "large",
    data = glob(["serve/**/*.yaml"]),
    env = {
        "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING": "1",
        "RAY_SERVE_ENABLE_HA_PROXY": "1",
    },
    files = [
        "serve/cpu/deployments/data_parallel/test_dp_direct_streaming.py",
        "serve/cpu/deployments/prefill_decode_disagg/test_pd_direct_streaming.py",
        "serve/cpu/deployments/routers/test_router_direct_streaming.py",
    ],
    tags = [
        "cpu",
        "exclusive",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)

# Multi-turn OpenAI-compatible HTTP benchmark (CLI + engine unit tests)
py_test_module_list(
    size = "large",
    files = glob(["serve/benchmark/test_*.py"]),
    tags = [
        "cpu",
        "exclusive",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)

# Large GPU tests
py_test_module_list(
    size = "large",
    data = glob(["serve/**/*.yaml"]),
    files = glob(["serve/gpu/**/test_*.py"]),
    tags = [
        "exclusive",
        "gpu",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)

py_test_module_list(
    size = "large",
    data = glob(["serve/**/*.yaml"]),
    env = {
        "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING": "1",
        # Direct streaming wires LLMRouter as the ingress_request_router,
        # which build_app.py rejects unless HAProxy is enabled.
        "RAY_SERVE_ENABLE_HA_PROXY": "1",
    },
    files = ["serve/gpu/integration/test_openai_compatibility.py"],
    name_suffix = "_direct_streaming",
    tags = [
        "exclusive",
        "gpu",
        "team:llm",
    ],
    deps = [
        ":conftest",
        "//:ray_lib",
    ],
)