backendai_env_test.rs

// Copyright 2025 Lablup Inc. and Jeongkyu Shin
//
// 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.

use bssh::config::Config;
use once_cell::sync::Lazy;
use std::env;
use tokio::sync::Mutex;

// Global mutex to serialize tests that modify environment variables
static ENV_MUTEX: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));

#[tokio::test]
async fn test_backendai_env_auto_detection() {
    let _guard = ENV_MUTEX.lock().await;

    // Save original env vars
    let orig_hosts = env::var("BACKENDAI_CLUSTER_HOSTS").ok();
    let orig_host = env::var("BACKENDAI_CLUSTER_HOST").ok();
    let orig_role = env::var("BACKENDAI_CLUSTER_ROLE").ok();

    // Set Backend.AI environment variables
    unsafe {
        env::set_var("BACKENDAI_CLUSTER_HOSTS", "node1.ai,node2.ai,node3.ai");
        env::set_var("BACKENDAI_CLUSTER_HOST", "node1.ai");
        env::set_var("BACKENDAI_CLUSTER_ROLE", "main");
    }

    // Create a temporary directory for the test
    let temp_dir = tempfile::tempdir().unwrap();
    let nonexistent_path = temp_dir.path().join("nonexistent.yaml");

    // Load config with priority (should detect Backend.AI env)
    let config = Config::load_with_priority(&nonexistent_path)
        .await
        .expect("Config should load with Backend.AI env");

    // Check that bai_auto cluster was created
    assert!(config.clusters.contains_key("bai_auto"));

    // Get the bai_auto cluster
    let cluster = config.clusters.get("bai_auto").unwrap();

    // Verify SSH key is set to Backend.AI cluster key
    assert_eq!(
        cluster.defaults.ssh_key,
        Some("/home/config/ssh/id_cluster".to_string()),
        "Backend.AI cluster should use /home/config/ssh/id_cluster as SSH key"
    );

    // Verify nodes were parsed correctly
    assert_eq!(cluster.nodes.len(), 3);

    // Resolve nodes for the bai_auto cluster
    let nodes = config
        .resolve_nodes("bai_auto")
        .expect("Should resolve bai_auto nodes");
    assert_eq!(nodes.len(), 3);

    // Check node details
    assert_eq!(nodes[0].host, "node1.ai");
    assert_eq!(nodes[0].port, 2200); // Backend.AI default port
    assert_eq!(nodes[1].host, "node2.ai");
    assert_eq!(nodes[2].host, "node3.ai");

    // Verify get_ssh_key returns the correct key for Backend.AI cluster
    assert_eq!(
        config.get_ssh_key(Some("bai_auto")),
        Some("/home/config/ssh/id_cluster".to_string()),
        "get_ssh_key should return Backend.AI cluster key path"
    );

    // Restore original env vars
    unsafe {
        if let Some(val) = orig_hosts {
            env::set_var("BACKENDAI_CLUSTER_HOSTS", val);
        } else {
            env::remove_var("BACKENDAI_CLUSTER_HOSTS");
        }

        if let Some(val) = orig_host {
            env::set_var("BACKENDAI_CLUSTER_HOST", val);
        } else {
            env::remove_var("BACKENDAI_CLUSTER_HOST");
        }

        if let Some(val) = orig_role {
            env::set_var("BACKENDAI_CLUSTER_ROLE", val);
        } else {
            env::remove_var("BACKENDAI_CLUSTER_ROLE");
        }
    }
}

#[tokio::test]
async fn test_backendai_env_with_single_host() {
    let _guard = ENV_MUTEX.lock().await;

    // Save original env vars
    let orig_hosts = env::var("BACKENDAI_CLUSTER_HOSTS").ok();
    let orig_host = env::var("BACKENDAI_CLUSTER_HOST").ok();
    let orig_role = env::var("BACKENDAI_CLUSTER_ROLE").ok();

    // Set Backend.AI environment variables with single host
    unsafe {
        env::set_var("BACKENDAI_CLUSTER_HOSTS", "single-node.ai");
        env::set_var("BACKENDAI_CLUSTER_HOST", "single-node.ai");
        // Explicitly remove ROLE to avoid contamination from previous tests
        env::remove_var("BACKENDAI_CLUSTER_ROLE");
    }

    // Create a temporary directory for the test
    let temp_dir = tempfile::tempdir().unwrap();
    let nonexistent_path = temp_dir.path().join("nonexistent.yaml");

    // Load config
    let config = Config::load_with_priority(&nonexistent_path)
        .await
        .expect("Config should load");

    // Verify bai_auto cluster exists
    assert!(config.clusters.contains_key("bai_auto"));

    let nodes = config
        .resolve_nodes("bai_auto")
        .expect("Should resolve nodes");
    assert_eq!(nodes.len(), 1);
    assert_eq!(nodes[0].host, "single-node.ai");
    assert_eq!(nodes[0].port, 2200);

    // Restore
    unsafe {
        if let Some(val) = orig_hosts {
            env::set_var("BACKENDAI_CLUSTER_HOSTS", val);
        } else {
            env::remove_var("BACKENDAI_CLUSTER_HOSTS");
        }

        if let Some(val) = orig_host {
            env::set_var("BACKENDAI_CLUSTER_HOST", val);
        } else {
            env::remove_var("BACKENDAI_CLUSTER_HOST");
        }

        if let Some(val) = orig_role {
            env::set_var("BACKENDAI_CLUSTER_ROLE", val);
        } else {
            env::remove_var("BACKENDAI_CLUSTER_ROLE");
        }
    }
}

#[tokio::test]
async fn test_no_backendai_env() {
    let _guard = ENV_MUTEX.lock().await;

    // Save and clear Backend.AI env vars
    let orig_hosts = env::var("BACKENDAI_CLUSTER_HOSTS").ok();
    let orig_host = env::var("BACKENDAI_CLUSTER_HOST").ok();

    unsafe {
        env::remove_var("BACKENDAI_CLUSTER_HOSTS");
        env::remove_var("BACKENDAI_CLUSTER_HOST");
        env::remove_var("BACKENDAI_CLUSTER_ROLE");
    }

    // Load config without Backend.AI env
    let config = Config::load_with_priority(&std::path::PathBuf::from("nonexistent.yaml"))
        .await
        .expect("Config should load");

    // Verify no backendai cluster was created
    assert!(!config.clusters.contains_key("backendai"));

    // Restore if needed
    unsafe {
        if let Some(val) = orig_hosts {
            env::set_var("BACKENDAI_CLUSTER_HOSTS", val);
        }
        if let Some(val) = orig_host {
            env::set_var("BACKENDAI_CLUSTER_HOST", val);
        }
    }
}