"""Tests for multi-node execution support."""

import os
from unittest.mock import patch

from bilgeai.cli.commands.run import _get_master_addr, _get_master_port
from bilgeai.optimizer.models import TrainingConfig


class TestMasterAddr:
    def test_default_localhost(self):
        with patch.dict(os.environ, {}, clear=True):
            assert _get_master_addr() == "localhost"

    def test_from_env(self):
        with patch.dict(os.environ, {"MASTER_ADDR": "192.168.1.100"}):
            assert _get_master_addr() == "192.168.1.100"


class TestMasterPort:
    def test_default_port(self):
        with patch.dict(os.environ, {}, clear=True):
            assert _get_master_port() == "29500"

    def test_from_env(self):
        with patch.dict(os.environ, {"MASTER_PORT": "12345"}):
            assert _get_master_port() == "12345"


class TestMultiNodeConfig:
    def test_default_single_node(self):
        config = TrainingConfig()
        assert config.num_nodes == 1

    def test_multi_node_setting(self):
        config = TrainingConfig(num_nodes=4, num_gpus=8)
        assert config.num_nodes == 4
        assert config.num_gpus == 8

    def test_multi_node_with_distributed(self):
        config = TrainingConfig(
            num_nodes=2,
            num_gpus=4,
            distributed_strategy="fsdp",
        )
        assert config.distributed_strategy == "fsdp"
        assert config.num_nodes == 2
