task_pilot/backend/tests/uv.rs

use std::{fs, path::Path};

use axum::{
    body::Body,
    http::{Method, Request, StatusCode, header},
};
use http_body_util::BodyExt;
use serde_json::{Value, json};
use tempfile::TempDir;
use tower::ServiceExt;

async fn test_app() -> (axum::Router, TempDir) {
    let temp_dir = TempDir::new().expect("temporary directory should be created");
    let pool = task_pilot_backend::database::connect("sqlite::memory:")
        .await
        .expect("test database should initialize");
    let app = task_pilot_backend::app(pool, temp_dir.path().join("logs"))
        .await
        .expect("test app should start");
    (app, temp_dir)
}

async fn auth_token(app: &axum::Router) -> String {
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/v1/session")
                .body(Body::empty())
                .expect("request should build"),
        )
        .await
        .expect("router should respond");
    let body = response
        .into_body()
        .collect()
        .await
        .expect("body should be readable")
        .to_bytes();
    let json: Value = serde_json::from_slice(&body).expect("session should be JSON");
    json["token"]
        .as_str()
        .expect("token should be a string")
        .to_owned()
}

async fn send_json(
    app: &axum::Router,
    method: Method,
    uri: &str,
    token: Option<&str>,
    body: Option<Value>,
) -> (StatusCode, Value) {
    let mut builder = Request::builder().method(method).uri(uri);
    if let Some(token) = token {
        builder = builder.header("x-taskpilot-token", token);
    }
    let request = if let Some(body) = body {
        builder
            .header(header::CONTENT_TYPE, "application/json")
            .body(Body::from(
                serde_json::to_vec(&body).expect("body should serialize"),
            ))
            .expect("request should build")
    } else {
        builder.body(Body::empty()).expect("request should build")
    };
    let response = app
        .clone()
        .oneshot(request)
        .await
        .expect("router should respond");
    let status = response.status();
    let body = response
        .into_body()
        .collect()
        .await
        .expect("body should be readable")
        .to_bytes();

    (status, serde_json::from_slice(&body).unwrap_or(Value::Null))
}

#[tokio::test]
async fn uv_settings_default_to_the_managed_uv() {
    let (app, _temp_dir) = test_app().await;

    let (status, body) = send_json(&app, Method::GET, "/api/v1/uv", None, None).await;

    assert_eq!(status, StatusCode::OK);
    assert_eq!(body["source"], "managed");
    assert_eq!(body["custom_path"], Value::Null);
    let managed_dir = body["managed_dir"]
        .as_str()
        .expect("managed_dir should be a string");
    assert_eq!(
        Path::new(managed_dir).file_name().unwrap(),
        std::ffi::OsStr::new("uv")
    );
    // 不再读系统 PATH，候选里只有程序目录 uv（自定义路径未设置时也不出现）。
    let sources = body["candidates"]
        .as_array()
        .expect("candidates should be an array")
        .iter()
        .map(|candidate| candidate["source"].as_str().unwrap().to_owned())
        .collect::<Vec<_>>();
    assert_eq!(sources, vec!["managed"]);
}

#[tokio::test]
async fn uv_settings_reject_unusable_custom_path_and_persist_valid_ones() {
    let (app, temp_dir) = test_app().await;
    let token = auth_token(&app).await;

    let (status, body) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        Some(&token),
        Some(json!({ "source": "custom", "custom_path": null })),
    )
    .await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
    assert_eq!(body["field"], "custom_path");

    let (status, body) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        Some(&token),
        Some(json!({ "source": "custom", "custom_path": "C:\\missing\\uv.exe" })),
    )
    .await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
    assert_eq!(body["field"], "custom_path");

    // 文件存在但不是可运行的 uv，也要拒绝保存。
    let fake = temp_dir.path().join("fake.exe");
    fs::write(&fake, b"not a program").expect("fake executable should be written");
    let (status, body) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        Some(&token),
        Some(json!({
            "source": "custom",
            "custom_path": fake.to_string_lossy()
        })),
    )
    .await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
    assert_eq!(body["field"], "custom_path");

    // 切回程序目录 uv 时保留已填写的自定义路径，方便再次切换。
    let (status, saved) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        Some(&token),
        Some(json!({
            "source": "managed",
            "custom_path": "  C:\\tools\\uv.exe  "
        })),
    )
    .await;
    assert_eq!(status, StatusCode::OK);
    assert_eq!(saved["source"], "managed");
    assert_eq!(saved["custom_path"], "C:\\tools\\uv.exe");

    let (status, body) = send_json(&app, Method::GET, "/api/v1/uv", None, None).await;
    assert_eq!(status, StatusCode::OK);
    assert_eq!(body["source"], "managed");
    assert_eq!(body["custom_path"], "C:\\tools\\uv.exe");
    assert!(
        body["candidates"]
            .as_array()
            .expect("candidates should be an array")
            .iter()
            .any(|candidate| candidate["source"] == "custom")
    );
}

#[tokio::test]
async fn uv_settings_mutations_require_token() {
    let (app, _temp_dir) = test_app().await;

    let (status, body) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        None,
        Some(json!({ "source": "managed", "custom_path": null })),
    )
    .await;

    assert_eq!(status, StatusCode::FORBIDDEN);
    assert_eq!(body["code"], "FORBIDDEN");
}

#[tokio::test]
async fn uv_mirror_is_saved_with_the_settings_and_drives_both_urls() {
    let (app, _temp_dir) = test_app().await;
    let token = auth_token(&app).await;

    let (status, body) = send_json(&app, Method::GET, "/api/v1/uv", None, None).await;
    assert_eq!(status, StatusCode::OK);
    assert_eq!(
        body["mirror_url"],
        "https://mirrors.ustc.edu.cn/github-release"
    );
    assert_eq!(
        body["uv_download_url"],
        "https://mirrors.ustc.edu.cn/github-release/astral-sh/uv/LatestRelease"
    );
    assert_eq!(
        body["python_mirror_url"],
        "https://mirrors.ustc.edu.cn/github-release/astral-sh/python-build-standalone/releases/download"
    );
    // 解释器镜像作为默认环境变量出现在列表里，可以被显式覆盖
    assert_eq!(
        body["env_defaults"]["UV_PYTHON_INSTALL_MIRROR"],
        "https://mirrors.ustc.edu.cn/github-release/astral-sh/python-build-standalone/releases/download"
    );

    // 非法镜像地址会被拒绝
    let (status, body) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        Some(&token),
        Some(json!({
            "source": "managed",
            "mirror_url": "ftp://mirror.example"
        })),
    )
    .await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
    assert_eq!(body["field"], "url");

    // 保存自定义镜像：两个下载地址都由它派生
    let (status, saved) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        Some(&token),
        Some(json!({
            "source": "managed",
            "mirror_url": "https://mirror.example/github-release/"
        })),
    )
    .await;
    assert_eq!(status, StatusCode::OK);
    assert_eq!(saved["mirror_url"], "https://mirror.example/github-release");
    assert_eq!(
        saved["uv_download_url"],
        "https://mirror.example/github-release/astral-sh/uv/LatestRelease"
    );

    // 传 null 恢复内置默认
    let (status, restored) = send_json(
        &app,
        Method::PUT,
        "/api/v1/uv",
        Some(&token),
        Some(json!({ "source": "managed", "mirror_url": null })),
    )
    .await;
    assert_eq!(status, StatusCode::OK);
    assert_eq!(
        restored["mirror_url"],
        "https://mirrors.ustc.edu.cn/github-release"
    );
}

#[tokio::test]
async fn uv_download_requires_token() {
    let (app, _temp_dir) = test_app().await;

    let (status, body) = send_json(&app, Method::POST, "/api/v1/uv/download", None, None).await;

    assert_eq!(status, StatusCode::FORBIDDEN);
    assert_eq!(body["code"], "FORBIDDEN");
}