aws-bootstrap-g4dn 0.1.0__py3-none-any.whl

aws_bootstrap/tests/test_ec2.py

@@ -0,0 +1,313 @@
+"""Tests for EC2 helper functions."""
+
+from __future__ import annotations
+import io
+from datetime import UTC, datetime
+from unittest.mock import MagicMock
+
+import botocore.exceptions
+import click
+import pytest
+
+from aws_bootstrap.config import LaunchConfig
+from aws_bootstrap.ec2 import (
+    CLIError,
+    find_tagged_instances,
+    get_latest_ami,
+    get_spot_price,
+    launch_instance,
+    list_amis,
+    list_instance_types,
+    terminate_tagged_instances,
+)
+
+
+def test_cli_error_is_click_exception():
+    err = CLIError("something went wrong")
+    assert isinstance(err, click.ClickException)
+    assert err.format_message() == "something went wrong"
+
+
+def test_cli_error_show_outputs_red():
+    err = CLIError("bad input")
+    buf = io.StringIO()
+    err.show(file=buf)
+    output = buf.getvalue()
+    assert "Error: bad input" in output
+
+
+def test_get_latest_ami_picks_newest():
+    ec2 = MagicMock()
+    ec2.describe_images.return_value = {
+        "Images": [
+            {"ImageId": "ami-old", "Name": "DL AMI old", "CreationDate": "2024-01-01T00:00:00Z"},
+            {"ImageId": "ami-new", "Name": "DL AMI new", "CreationDate": "2025-06-01T00:00:00Z"},
+            {"ImageId": "ami-mid", "Name": "DL AMI mid", "CreationDate": "2025-01-01T00:00:00Z"},
+        ]
+    }
+    ami = get_latest_ami(ec2, "DL AMI*")
+    assert ami["ImageId"] == "ami-new"
+
+
+def test_get_latest_ami_no_results():
+    ec2 = MagicMock()
+    ec2.describe_images.return_value = {"Images": []}
+    with pytest.raises(click.ClickException, match="No AMI found"):
+        get_latest_ami(ec2, "nonexistent*")
+
+
+def _make_client_error(code: str, message: str = "test") -> botocore.exceptions.ClientError:
+    return botocore.exceptions.ClientError(
+        {"Error": {"Code": code, "Message": message}},
+        "RunInstances",
+    )
+
+
+def test_launch_instance_spot_quota_exceeded():
+    ec2 = MagicMock()
+    ec2.run_instances.side_effect = _make_client_error("MaxSpotInstanceCountExceeded")
+    config = LaunchConfig(spot=True)
+    with pytest.raises(click.ClickException, match="Spot instance quota exceeded"):
+        launch_instance(ec2, config, "ami-test", "sg-test")
+
+
+def test_launch_instance_vcpu_limit_exceeded():
+    ec2 = MagicMock()
+    ec2.run_instances.side_effect = _make_client_error("VcpuLimitExceeded")
+    config = LaunchConfig(spot=False)
+    with pytest.raises(click.ClickException, match="vCPU quota exceeded"):
+        launch_instance(ec2, config, "ami-test", "sg-test")
+
+
+def test_launch_instance_quota_error_includes_readme_hint():
+    ec2 = MagicMock()
+    ec2.run_instances.side_effect = _make_client_error("MaxSpotInstanceCountExceeded")
+    config = LaunchConfig(spot=True)
+    with pytest.raises(click.ClickException, match="README.md"):
+        launch_instance(ec2, config, "ami-test", "sg-test")
+
+
+def test_find_tagged_instances():
+    ec2 = MagicMock()
+    ec2.describe_instances.return_value = {
+        "Reservations": [
+            {
+                "Instances": [
+                    {
+                        "InstanceId": "i-abc123",
+                        "State": {"Name": "running"},
+                        "InstanceType": "g4dn.xlarge",
+                        "PublicIpAddress": "1.2.3.4",
+                        "LaunchTime": datetime(2025, 1, 1, tzinfo=UTC),
+                        "InstanceLifecycle": "spot",
+                        "Placement": {"AvailabilityZone": "us-west-2a"},
+                        "Tags": [
+                            {"Key": "Name", "Value": "aws-bootstrap-g4dn.xlarge"},
+                            {"Key": "created-by", "Value": "aws-bootstrap-g4dn"},
+                        ],
+                    }
+                ]
+            }
+        ]
+    }
+    instances = find_tagged_instances(ec2, "aws-bootstrap-g4dn")
+    assert len(instances) == 1
+    assert instances[0]["InstanceId"] == "i-abc123"
+    assert instances[0]["State"] == "running"
+    assert instances[0]["PublicIp"] == "1.2.3.4"
+    assert instances[0]["Name"] == "aws-bootstrap-g4dn.xlarge"
+    assert instances[0]["Lifecycle"] == "spot"
+    assert instances[0]["AvailabilityZone"] == "us-west-2a"
+
+
+def test_find_tagged_instances_on_demand_lifecycle():
+    """On-demand instances have no InstanceLifecycle key; should default to 'on-demand'."""
+    ec2 = MagicMock()
+    ec2.describe_instances.return_value = {
+        "Reservations": [
+            {
+                "Instances": [
+                    {
+                        "InstanceId": "i-ondemand",
+                        "State": {"Name": "running"},
+                        "InstanceType": "g4dn.xlarge",
+                        "PublicIpAddress": "5.6.7.8",
+                        "LaunchTime": datetime(2025, 1, 1, tzinfo=UTC),
+                        "Placement": {"AvailabilityZone": "us-west-2b"},
+                        "Tags": [
+                            {"Key": "Name", "Value": "aws-bootstrap-g4dn.xlarge"},
+                        ],
+                    }
+                ]
+            }
+        ]
+    }
+    instances = find_tagged_instances(ec2, "aws-bootstrap-g4dn")
+    assert len(instances) == 1
+    assert instances[0]["Lifecycle"] == "on-demand"
+    assert instances[0]["AvailabilityZone"] == "us-west-2b"
+
+
+def test_find_tagged_instances_empty():
+    ec2 = MagicMock()
+    ec2.describe_instances.return_value = {"Reservations": []}
+    assert find_tagged_instances(ec2, "aws-bootstrap-g4dn") == []
+
+
+def test_get_spot_price_returns_price():
+    ec2 = MagicMock()
+    ec2.describe_spot_price_history.return_value = {"SpotPriceHistory": [{"SpotPrice": "0.1578"}]}
+    price = get_spot_price(ec2, "g4dn.xlarge", "us-west-2a")
+    assert price == 0.1578
+    ec2.describe_spot_price_history.assert_called_once()
+
+
+def test_get_spot_price_returns_none_when_empty():
+    ec2 = MagicMock()
+    ec2.describe_spot_price_history.return_value = {"SpotPriceHistory": []}
+    price = get_spot_price(ec2, "g4dn.xlarge", "us-west-2a")
+    assert price is None
+
+
+def test_terminate_tagged_instances():
+    ec2 = MagicMock()
+    ec2.terminate_instances.return_value = {
+        "TerminatingInstances": [
+            {
+                "InstanceId": "i-abc123",
+                "PreviousState": {"Name": "running"},
+                "CurrentState": {"Name": "shutting-down"},
+            }
+        ]
+    }
+    changes = terminate_tagged_instances(ec2, ["i-abc123"])
+    assert len(changes) == 1
+    assert changes[0]["InstanceId"] == "i-abc123"
+    ec2.terminate_instances.assert_called_once_with(InstanceIds=["i-abc123"])
+
+
+# ---------------------------------------------------------------------------
+# list_instance_types
+# ---------------------------------------------------------------------------
+
+
+def test_list_instance_types_returns_sorted():
+    ec2 = MagicMock()
+    paginator = MagicMock()
+    ec2.get_paginator.return_value = paginator
+    paginator.paginate.return_value = [
+        {
+            "InstanceTypes": [
+                {
+                    "InstanceType": "g4dn.xlarge",
+                    "VCpuInfo": {"DefaultVCpus": 4},
+                    "MemoryInfo": {"SizeInMiB": 16384},
+                    "GpuInfo": {"Gpus": [{"Count": 1, "Name": "T4", "MemoryInfo": {"SizeInMiB": 16384}}]},
+                },
+                {
+                    "InstanceType": "g4dn.2xlarge",
+                    "VCpuInfo": {"DefaultVCpus": 8},
+                    "MemoryInfo": {"SizeInMiB": 32768},
+                    "GpuInfo": {"Gpus": [{"Count": 1, "Name": "T4", "MemoryInfo": {"SizeInMiB": 16384}}]},
+                },
+            ]
+        }
+    ]
+    results = list_instance_types(ec2, "g4dn")
+    assert len(results) == 2
+    # sorted by name — 2xlarge < xlarge lexicographically
+    assert results[0]["InstanceType"] == "g4dn.2xlarge"
+    assert results[1]["InstanceType"] == "g4dn.xlarge"
+    assert results[1]["VCpuCount"] == 4
+    assert results[1]["MemoryMiB"] == 16384
+    assert "T4" in results[1]["GpuSummary"]
+
+
+def test_list_instance_types_no_gpu():
+    ec2 = MagicMock()
+    paginator = MagicMock()
+    ec2.get_paginator.return_value = paginator
+    paginator.paginate.return_value = [
+        {
+            "InstanceTypes": [
+                {
+                    "InstanceType": "t3.medium",
+                    "VCpuInfo": {"DefaultVCpus": 2},
+                    "MemoryInfo": {"SizeInMiB": 4096},
+                },
+            ]
+        }
+    ]
+    results = list_instance_types(ec2, "t3")
+    assert len(results) == 1
+    assert results[0]["GpuSummary"] == ""
+
+
+def test_list_instance_types_empty():
+    ec2 = MagicMock()
+    paginator = MagicMock()
+    ec2.get_paginator.return_value = paginator
+    paginator.paginate.return_value = [{"InstanceTypes": []}]
+    results = list_instance_types(ec2, "nonexistent")
+    assert results == []
+
+
+# ---------------------------------------------------------------------------
+# list_amis
+# ---------------------------------------------------------------------------
+
+
+def test_list_amis_sorted_newest_first():
+    ec2 = MagicMock()
+    ec2.describe_images.return_value = {
+        "Images": [
+            {
+                "ImageId": "ami-old",
+                "Name": "DL AMI old",
+                "CreationDate": "2024-01-01T00:00:00Z",
+                "Architecture": "x86_64",
+            },
+            {
+                "ImageId": "ami-new",
+                "Name": "DL AMI new",
+                "CreationDate": "2025-06-01T00:00:00Z",
+                "Architecture": "x86_64",
+            },
+        ]
+    }
+    results = list_amis(ec2, "DL AMI*")
+    assert len(results) == 2
+    assert results[0]["ImageId"] == "ami-new"
+    assert results[1]["ImageId"] == "ami-old"
+
+
+def test_list_amis_empty():
+    ec2 = MagicMock()
+    ec2.describe_images.return_value = {"Images": []}
+    results = list_amis(ec2, "nonexistent*")
+    assert results == []
+
+
+def test_list_amis_limited_to_20():
+    ec2 = MagicMock()
+    ec2.describe_images.return_value = {
+        "Images": [
+            {
+                "ImageId": f"ami-{i:03d}",
+                "Name": f"AMI {i}",
+                "CreationDate": f"2025-01-{i + 1:02d}T00:00:00Z",
+                "Architecture": "x86_64",
+            }
+            for i in range(25)
+        ]
+    }
+    results = list_amis(ec2, "AMI*")
+    assert len(results) == 20
+
+
+def test_list_amis_uses_owner_hint_for_deep_learning():
+    ec2 = MagicMock()
+    ec2.describe_images.return_value = {"Images": []}
+    list_amis(ec2, "Deep Learning Base*")
+    call_kwargs = ec2.describe_images.call_args[1]
+    assert call_kwargs["Owners"] == ["amazon"]

aws_bootstrap/tests/test_ssh_config.py

@@ -0,0 +1,297 @@
+"""Tests for SSH config management (add/remove/find/list host stanzas)."""
+
+from __future__ import annotations
+import os
+import stat
+from pathlib import Path
+
+from aws_bootstrap.ssh import (
+    _next_alias,
+    _read_ssh_config,
+    add_ssh_host,
+    find_ssh_alias,
+    list_ssh_hosts,
+    remove_ssh_host,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+KEY_PATH = Path("/home/user/.ssh/id_ed25519.pub")
+
+
+def _config_path(tmp_path: Path) -> Path:
+    """Return a config path inside tmp_path (doesn't create the file)."""
+    return tmp_path / ".ssh" / "config"
+
+
+# ---------------------------------------------------------------------------
+# Stanza creation
+# ---------------------------------------------------------------------------
+
+
+def test_add_creates_stanza(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    content = cfg.read_text()
+    assert "# >>> aws-bootstrap [i-abc123] >>>" in content
+    assert "# <<< aws-bootstrap [i-abc123] <<<" in content
+    assert "Host aws-gpu1" in content
+    assert "HostName 1.2.3.4" in content
+    assert "User ubuntu" in content
+    assert "IdentityFile /home/user/.ssh/id_ed25519" in content
+
+
+def test_add_returns_alias(tmp_path):
+    cfg = _config_path(tmp_path)
+    alias = add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    assert alias == "aws-gpu1"
+
+
+def test_add_uses_private_key_path(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", Path("/keys/mykey.pub"), config_path=cfg)
+    content = cfg.read_text()
+    assert "IdentityFile /keys/mykey" in content
+    assert ".pub" not in content.split("IdentityFile")[1].split("\n")[0]
+
+
+def test_add_includes_user(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ec2-user", KEY_PATH, config_path=cfg)
+    content = cfg.read_text()
+    assert "User ec2-user" in content
+
+
+def test_add_includes_strict_host_checking(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    content = cfg.read_text()
+    assert "StrictHostKeyChecking no" in content
+    assert "UserKnownHostsFile /dev/null" in content
+
+
+# ---------------------------------------------------------------------------
+# Multiple instances
+# ---------------------------------------------------------------------------
+
+
+def test_second_host_gets_gpu2(tmp_path):
+    cfg = _config_path(tmp_path)
+    a1 = add_ssh_host("i-111", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    a2 = add_ssh_host("i-222", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    assert a1 == "aws-gpu1"
+    assert a2 == "aws-gpu2"
+
+
+def test_next_alias_empty():
+    assert _next_alias("") == "aws-gpu1"
+
+
+def test_next_alias_custom_prefix():
+    assert _next_alias("", prefix="dev-box") == "dev-box1"
+
+
+def test_next_alias_skips_user_hosts():
+    """User-defined hosts with similar names should be ignored."""
+    content = "Host aws-gpu99\n    HostName 1.2.3.4\n"
+    # No marker blocks, so this should be ignored
+    assert _next_alias(content) == "aws-gpu1"
+
+
+# ---------------------------------------------------------------------------
+# Preserving existing config
+# ---------------------------------------------------------------------------
+
+
+def test_preserves_existing_stanzas(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    existing = "Host myserver\n    HostName 10.0.0.1\n    User admin\n"
+    cfg.write_text(existing)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    content = cfg.read_text()
+    assert "Host myserver" in content
+    assert "HostName 10.0.0.1" in content
+    assert "Host aws-gpu1" in content
+
+
+def test_preserves_trailing_newline(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("Host foo\n    HostName bar\n")
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    content = cfg.read_text()
+    # Should not have triple+ blank lines
+    assert "\n\n\n" not in content
+
+
+def test_add_to_nonexistent_creates_dir_and_file(tmp_path):
+    cfg = tmp_path / "brand_new" / ".ssh" / "config"
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    assert cfg.exists()
+    assert cfg.parent.exists()
+
+
+# ---------------------------------------------------------------------------
+# Removal
+# ---------------------------------------------------------------------------
+
+
+def test_remove_returns_alias(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    removed = remove_ssh_host("i-abc123", config_path=cfg)
+    assert removed == "aws-gpu1"
+
+
+def test_remove_returns_none_when_missing(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("")
+    assert remove_ssh_host("i-missing", config_path=cfg) is None
+
+
+def test_remove_preserves_other_stanzas(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    remove_ssh_host("i-111", config_path=cfg)
+    content = cfg.read_text()
+    assert "i-111" not in content
+    assert "Host aws-gpu2" in content
+    assert "HostName 2.2.2.2" in content
+
+
+def test_remove_preserves_user_config(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("Host myserver\n    HostName 10.0.0.1\n")
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    remove_ssh_host("i-abc123", config_path=cfg)
+    content = cfg.read_text()
+    assert "Host myserver" in content
+    assert "HostName 10.0.0.1" in content
+    assert "aws-gpu" not in content
+
+
+def test_remove_cleans_trailing_blanks(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    remove_ssh_host("i-abc123", config_path=cfg)
+    content = cfg.read_text()
+    assert "\n\n\n" not in content
+
+
+# ---------------------------------------------------------------------------
+# Idempotency
+# ---------------------------------------------------------------------------
+
+
+def test_add_idempotent_updates_ip(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-abc123", "5.6.7.8", "ubuntu", KEY_PATH, config_path=cfg)
+    content = cfg.read_text()
+    assert "HostName 5.6.7.8" in content
+    assert "HostName 1.2.3.4" not in content
+
+
+def test_add_idempotent_preserves_alias(tmp_path):
+    cfg = _config_path(tmp_path)
+    a1 = add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    a2 = add_ssh_host("i-abc123", "5.6.7.8", "ubuntu", KEY_PATH, config_path=cfg)
+    assert a1 == a2 == "aws-gpu1"
+
+
+# ---------------------------------------------------------------------------
+# Lookup
+# ---------------------------------------------------------------------------
+
+
+def test_find_alias_returns_alias(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    assert find_ssh_alias("i-abc123", config_path=cfg) == "aws-gpu1"
+
+
+def test_find_alias_returns_none(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("")
+    assert find_ssh_alias("i-missing", config_path=cfg) is None
+
+
+def test_list_hosts_returns_all(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    hosts = list_ssh_hosts(config_path=cfg)
+    assert hosts == {"i-111": "aws-gpu1", "i-222": "aws-gpu2"}
+
+
+def test_list_hosts_ignores_user_stanzas(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("Host myserver\n    HostName 10.0.0.1\n")
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    hosts = list_ssh_hosts(config_path=cfg)
+    assert "myserver" not in hosts.values()
+    assert hosts == {"i-abc123": "aws-gpu1"}
+
+
+# ---------------------------------------------------------------------------
+# Safety / edge cases
+# ---------------------------------------------------------------------------
+
+
+def test_file_permissions_0600(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    mode = stat.S_IMODE(os.stat(cfg).st_mode)
+    assert mode == 0o600
+
+
+def test_dir_permissions_0700(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    mode = stat.S_IMODE(os.stat(cfg.parent).st_mode)
+    assert mode == 0o700
+
+
+def test_handles_empty_file(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("")
+    alias = add_ssh_host("i-abc123", "1.2.3.4", "ubuntu", KEY_PATH, config_path=cfg)
+    assert alias == "aws-gpu1"
+    assert "Host aws-gpu1" in cfg.read_text()
+
+
+def test_malformed_marker_left_alone(tmp_path):
+    """Orphaned begin marker without end marker should not cause deletion."""
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    orphaned = "# >>> aws-bootstrap [i-orphan] >>>\nHost aws-gpu99\n    HostName 9.9.9.9\n"
+    cfg.write_text(orphaned)
+    removed = remove_ssh_host("i-orphan", config_path=cfg)
+    assert removed is None
+    content = cfg.read_text()
+    assert "aws-gpu99" in content
+
+
+def test_read_nonexistent_returns_empty(tmp_path):
+    cfg = tmp_path / "does_not_exist"
+    assert _read_ssh_config(cfg) == ""
+
+
+def test_list_hosts_nonexistent_file(tmp_path):
+    cfg = tmp_path / "no_such_file"
+    assert list_ssh_hosts(config_path=cfg) == {}
+
+
+def test_remove_nonexistent_file(tmp_path):
+    cfg = tmp_path / "no_such_file"
+    assert remove_ssh_host("i-abc123", config_path=cfg) is None