"""Tests for cloud training command generators."""

import pytest

from orbit.training.cloud import CloudTrainingConfig, generate_cloud_command


class TestCloudCommands:
    @pytest.mark.parametrize("cloud", ["gcp", "aws", "runpod", "lambda"])
    def test_generates_valid_config(self, cloud):
        config = generate_cloud_command(
            dataset_repo="test/dataset",
            policy_type="act",
            cloud=cloud,
            gpu_type="a100",
        )
        assert isinstance(config, CloudTrainingConfig)
        assert config.cloud == cloud
        assert len(config.command) > 10
        assert config.estimated_cost
        assert len(config.notes) > 0

    def test_unknown_cloud_raises(self):
        with pytest.raises(ValueError, match="Unknown cloud"):
            generate_cloud_command(
                dataset_repo="test/dataset",
                policy_type="act",
                cloud="nonexistent",
            )


class TestGCPCommand:
    def test_gcloud_run_jobs(self):
        config = generate_cloud_command("test/data", "act", "gcp", "a100")
        assert "gcloud run jobs" in config.command
        assert "nvidia-a100" in config.command

    def test_l4_gpu(self):
        config = generate_cloud_command("test/data", "act", "gcp", "l4")
        assert "nvidia-l4" in config.command


class TestAWSCommand:
    def test_sagemaker_config(self):
        config = generate_cloud_command("test/data", "act", "aws", "a100")
        assert "PyTorch" in config.command
        assert "ml.p4d" in config.command

    def test_g5_instance(self):
        config = generate_cloud_command("test/data", "act", "aws", "l4")
        assert "ml.g5" in config.command


class TestRunPodCommand:
    def test_runpodctl_create(self):
        config = generate_cloud_command("test/data", "act", "runpod", "a100")
        assert "runpodctl create pod" in config.command
        assert "NVIDIA A100" in config.command


class TestLambdaCommand:
    def test_lambda_cloud_launch(self):
        config = generate_cloud_command("test/data", "act", "lambda", "a100")
        assert "lambda-cloud" in config.command or "gpu_1x_a100" in config.command
