bssh 2.1.2

Parallel SSH command execution tool for cluster management
Documentation
// 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.

mod common;

use bssh::config::Config;
use common::EnvGuard;
use serial_test::serial;

#[tokio::test]
#[serial]
async fn test_backendai_env_auto_detection() {
    // Guards restore prior values on drop.
    let _hosts = EnvGuard::set("BACKENDAI_CLUSTER_HOSTS", "node1.ai,node2.ai,node3.ai");
    let _host = EnvGuard::set("BACKENDAI_CLUSTER_HOST", "node1.ai");
    let _role = EnvGuard::set("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"
    );
}

#[tokio::test]
#[serial]
async fn test_backendai_env_with_single_host() {
    // Guards restore prior values on drop.
    let _hosts = EnvGuard::set("BACKENDAI_CLUSTER_HOSTS", "single-node.ai");
    let _host = EnvGuard::set("BACKENDAI_CLUSTER_HOST", "single-node.ai");
    // Explicitly remove ROLE to avoid contamination from previous tests
    let _role = EnvGuard::remove("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);
}

#[tokio::test]
#[serial]
async fn test_no_backendai_env() {
    // Clear all Backend.AI env vars; guards restore prior values on drop.
    let _hosts = EnvGuard::remove("BACKENDAI_CLUSTER_HOSTS");
    let _host = EnvGuard::remove("BACKENDAI_CLUSTER_HOST");
    let _role = EnvGuard::remove("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"));
}