aws-bootstrap-g4dn 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

aws_bootstrap/tests/test_config.py

@@ -20,6 +20,12 @@ def test_defaults():
     assert config.dry_run is False
 
 
+def test_ebs_fields_default_none():
+    config = LaunchConfig()
+    assert config.ebs_storage is None
+    assert config.ebs_volume_id is None
+
+
 def test_overrides():
     config = LaunchConfig(
         instance_type="g5.xlarge",
@@ -33,3 +39,15 @@ def test_overrides():
     assert config.spot is False
     assert config.volume_size == 200
     assert config.key_path == Path("/tmp/test.pub")
+
+
+def test_ebs_storage_override():
+    config = LaunchConfig(ebs_storage=96)
+    assert config.ebs_storage == 96
+    assert config.ebs_volume_id is None
+
+
+def test_ebs_volume_id_override():
+    config = LaunchConfig(ebs_volume_id="vol-abc123")
+    assert config.ebs_volume_id == "vol-abc123"
+    assert config.ebs_storage is None

aws_bootstrap/tests/test_ebs.py

@@ -0,0 +1,245 @@
+"""Tests for EBS data volume operations in ec2.py."""
+
+from __future__ import annotations
+from unittest.mock import MagicMock
+
+import botocore.exceptions
+import pytest
+
+from aws_bootstrap.ec2 import (
+    EBS_DEVICE_NAME,
+    CLIError,
+    attach_ebs_volume,
+    create_ebs_volume,
+    delete_ebs_volume,
+    detach_ebs_volume,
+    find_ebs_volumes_for_instance,
+    validate_ebs_volume,
+)
+
+
+# ---------------------------------------------------------------------------
+# create_ebs_volume
+# ---------------------------------------------------------------------------
+
+
+def test_create_ebs_volume():
+    ec2 = MagicMock()
+    ec2.create_volume.return_value = {"VolumeId": "vol-abc123"}
+    waiter = MagicMock()
+    ec2.get_waiter.return_value = waiter
+
+    vol_id = create_ebs_volume(ec2, 96, "us-west-2a", "aws-bootstrap-g4dn", "i-test123")
+
+    assert vol_id == "vol-abc123"
+    ec2.create_volume.assert_called_once()
+    create_kwargs = ec2.create_volume.call_args[1]
+    assert create_kwargs["AvailabilityZone"] == "us-west-2a"
+    assert create_kwargs["Size"] == 96
+    assert create_kwargs["VolumeType"] == "gp3"
+
+    # Check tags
+    tags = create_kwargs["TagSpecifications"][0]["Tags"]
+    tag_dict = {t["Key"]: t["Value"] for t in tags}
+    assert tag_dict["created-by"] == "aws-bootstrap-g4dn"
+    assert tag_dict["Name"] == "aws-bootstrap-data-i-test123"
+    assert tag_dict["aws-bootstrap-instance"] == "i-test123"
+
+    ec2.get_waiter.assert_called_once_with("volume_available")
+    waiter.wait.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# validate_ebs_volume
+# ---------------------------------------------------------------------------
+
+
+def test_validate_ebs_volume_valid():
+    ec2 = MagicMock()
+    ec2.describe_volumes.return_value = {
+        "Volumes": [
+            {
+                "VolumeId": "vol-abc123",
+                "State": "available",
+                "AvailabilityZone": "us-west-2a",
+                "Size": 100,
+            }
+        ]
+    }
+    vol = validate_ebs_volume(ec2, "vol-abc123", "us-west-2a")
+    assert vol["VolumeId"] == "vol-abc123"
+
+
+def test_validate_ebs_volume_wrong_az():
+    ec2 = MagicMock()
+    ec2.describe_volumes.return_value = {
+        "Volumes": [
+            {
+                "VolumeId": "vol-abc123",
+                "State": "available",
+                "AvailabilityZone": "us-east-1a",
+                "Size": 100,
+            }
+        ]
+    }
+    with pytest.raises(CLIError, match="us-east-1a"):
+        validate_ebs_volume(ec2, "vol-abc123", "us-west-2a")
+
+
+def test_validate_ebs_volume_in_use():
+    ec2 = MagicMock()
+    ec2.describe_volumes.return_value = {
+        "Volumes": [
+            {
+                "VolumeId": "vol-abc123",
+                "State": "in-use",
+                "AvailabilityZone": "us-west-2a",
+                "Size": 100,
+            }
+        ]
+    }
+    with pytest.raises(CLIError, match="in-use"):
+        validate_ebs_volume(ec2, "vol-abc123", "us-west-2a")
+
+
+def test_validate_ebs_volume_not_found():
+    ec2 = MagicMock()
+    ec2.describe_volumes.side_effect = botocore.exceptions.ClientError(
+        {"Error": {"Code": "InvalidVolume.NotFound", "Message": "not found"}},
+        "DescribeVolumes",
+    )
+    with pytest.raises(CLIError, match="not found"):
+        validate_ebs_volume(ec2, "vol-notfound", "us-west-2a")
+
+
+def test_validate_ebs_volume_empty_response():
+    ec2 = MagicMock()
+    ec2.describe_volumes.return_value = {"Volumes": []}
+    with pytest.raises(CLIError, match="not found"):
+        validate_ebs_volume(ec2, "vol-empty", "us-west-2a")
+
+
+# ---------------------------------------------------------------------------
+# attach_ebs_volume
+# ---------------------------------------------------------------------------
+
+
+def test_attach_ebs_volume():
+    ec2 = MagicMock()
+    waiter = MagicMock()
+    ec2.get_waiter.return_value = waiter
+
+    attach_ebs_volume(ec2, "vol-abc123", "i-test123")
+
+    ec2.attach_volume.assert_called_once_with(
+        VolumeId="vol-abc123",
+        InstanceId="i-test123",
+        Device=EBS_DEVICE_NAME,
+    )
+    ec2.get_waiter.assert_called_once_with("volume_in_use")
+    waiter.wait.assert_called_once()
+
+
+def test_attach_ebs_volume_custom_device():
+    ec2 = MagicMock()
+    waiter = MagicMock()
+    ec2.get_waiter.return_value = waiter
+
+    attach_ebs_volume(ec2, "vol-abc123", "i-test123", device_name="/dev/sdg")
+
+    ec2.attach_volume.assert_called_once_with(
+        VolumeId="vol-abc123",
+        InstanceId="i-test123",
+        Device="/dev/sdg",
+    )
+
+
+# ---------------------------------------------------------------------------
+# detach_ebs_volume
+# ---------------------------------------------------------------------------
+
+
+def test_detach_ebs_volume():
+    ec2 = MagicMock()
+    waiter = MagicMock()
+    ec2.get_waiter.return_value = waiter
+
+    detach_ebs_volume(ec2, "vol-abc123")
+
+    ec2.detach_volume.assert_called_once_with(VolumeId="vol-abc123")
+    ec2.get_waiter.assert_called_once_with("volume_available")
+    waiter.wait.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# delete_ebs_volume
+# ---------------------------------------------------------------------------
+
+
+def test_delete_ebs_volume():
+    ec2 = MagicMock()
+    delete_ebs_volume(ec2, "vol-abc123")
+    ec2.delete_volume.assert_called_once_with(VolumeId="vol-abc123")
+
+
+# ---------------------------------------------------------------------------
+# find_ebs_volumes_for_instance
+# ---------------------------------------------------------------------------
+
+
+def test_find_ebs_volumes_for_instance():
+    ec2 = MagicMock()
+    ec2.describe_volumes.return_value = {
+        "Volumes": [
+            {
+                "VolumeId": "vol-data1",
+                "Size": 96,
+                "State": "in-use",
+                "Attachments": [{"Device": "/dev/sdf", "InstanceId": "i-test123"}],
+            }
+        ]
+    }
+    volumes = find_ebs_volumes_for_instance(ec2, "i-test123", "aws-bootstrap-g4dn")
+    assert len(volumes) == 1
+    assert volumes[0]["VolumeId"] == "vol-data1"
+    assert volumes[0]["Size"] == 96
+    assert volumes[0]["Device"] == "/dev/sdf"
+    assert volumes[0]["State"] == "in-use"
+
+
+def test_find_ebs_volumes_empty():
+    ec2 = MagicMock()
+    ec2.describe_volumes.return_value = {"Volumes": []}
+    volumes = find_ebs_volumes_for_instance(ec2, "i-test123", "aws-bootstrap-g4dn")
+    assert volumes == []
+
+
+def test_find_ebs_volumes_includes_available():
+    """Detached (available) volumes are still discovered by tags."""
+    ec2 = MagicMock()
+    ec2.describe_volumes.return_value = {
+        "Volumes": [
+            {
+                "VolumeId": "vol-avail",
+                "Size": 50,
+                "State": "available",
+                "Attachments": [],
+            }
+        ]
+    }
+    volumes = find_ebs_volumes_for_instance(ec2, "i-old", "aws-bootstrap-g4dn")
+    assert len(volumes) == 1
+    assert volumes[0]["VolumeId"] == "vol-avail"
+    assert volumes[0]["State"] == "available"
+    assert volumes[0]["Device"] == ""
+
+
+def test_find_ebs_volumes_client_error_returns_empty():
+    """ClientError (e.g. permissions) returns empty list instead of raising."""
+    ec2 = MagicMock()
+    ec2.describe_volumes.side_effect = botocore.exceptions.ClientError(
+        {"Error": {"Code": "UnauthorizedOperation", "Message": "no access"}},
+        "DescribeVolumes",
+    )
+    volumes = find_ebs_volumes_for_instance(ec2, "i-test", "aws-bootstrap-g4dn")
+    assert volumes == []

aws_bootstrap/tests/test_output.py

@@ -0,0 +1,192 @@
+"""Tests for the output formatting module."""
+
+from __future__ import annotations
+import json
+from datetime import UTC, datetime
+from pathlib import Path
+
+import click
+import yaml
+from click.testing import CliRunner
+
+from aws_bootstrap.output import OutputFormat, echo, emit, is_text
+
+
+def test_output_format_enum_values():
+    assert OutputFormat.TEXT.value == "text"
+    assert OutputFormat.JSON.value == "json"
+    assert OutputFormat.YAML.value == "yaml"
+    assert OutputFormat.TABLE.value == "table"
+
+
+def test_serialize_datetime():
+    """datetime objects should serialize to ISO format strings."""
+    dt = datetime(2025, 6, 15, 12, 30, 0, tzinfo=UTC)
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.JSON
+        emit({"timestamp": dt}, ctx=ctx)
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    data = json.loads(result.output)
+    assert data["timestamp"] == "2025-06-15T12:30:00+00:00"
+
+
+def test_serialize_path():
+    """Path objects should serialize to strings."""
+    p = Path("/home/user/.ssh/id_ed25519")
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.JSON
+        emit({"path": p}, ctx=ctx)
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    data = json.loads(result.output)
+    assert data["path"] == "/home/user/.ssh/id_ed25519"
+
+
+def test_emit_json():
+    """emit() should produce valid JSON in JSON mode."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.JSON
+        emit({"key": "value", "count": 42}, ctx=ctx)
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    data = json.loads(result.output)
+    assert data == {"key": "value", "count": 42}
+
+
+def test_emit_yaml():
+    """emit() should produce valid YAML in YAML mode."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.YAML
+        emit({"key": "value", "count": 42}, ctx=ctx)
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    data = yaml.safe_load(result.output)
+    assert data == {"key": "value", "count": 42}
+
+
+def test_emit_table_list():
+    """emit() should render a list of dicts as a table with headers."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.TABLE
+        emit(
+            [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
+            headers={"name": "Name", "age": "Age"},
+            ctx=ctx,
+        )
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    assert "Name" in result.output
+    assert "Age" in result.output
+    assert "Alice" in result.output
+    assert "Bob" in result.output
+
+
+def test_emit_table_dict():
+    """emit() should render a single dict as key-value pairs."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.TABLE
+        emit({"instance_id": "i-abc123", "state": "running"}, ctx=ctx)
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    assert "instance_id" in result.output
+    assert "i-abc123" in result.output
+    assert "running" in result.output
+
+
+def test_echo_suppressed_in_json_mode():
+    """echo() should produce no output when format is JSON."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.JSON
+        echo("This should not appear")
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    assert result.output == ""
+
+
+def test_echo_emits_in_text_mode():
+    """echo() should work normally in text mode."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.TEXT
+        echo("Hello world")
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+    assert "Hello world" in result.output
+
+
+def test_is_text_default():
+    """is_text() should return True when no context is set (default behavior)."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.TEXT
+        assert is_text(ctx) is True
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0
+
+
+def test_is_text_false_for_json():
+    """is_text() should return False when format is JSON."""
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.ensure_object(dict)
+        ctx.obj["output_format"] = OutputFormat.JSON
+        assert is_text(ctx) is False
+
+    runner = CliRunner()
+    result = runner.invoke(cli, [])
+    assert result.exit_code == 0

aws_bootstrap/tests/test_ssh_config.py

@@ -10,7 +10,9 @@ from aws_bootstrap.ssh import (
     _next_alias,
     _read_ssh_config,
     add_ssh_host,
+    cleanup_stale_ssh_hosts,
     find_ssh_alias,
+    find_stale_ssh_hosts,
     get_ssh_host_details,
     list_ssh_hosts,
     remove_ssh_host,
@@ -407,3 +409,77 @@ def test_resolve_unknown_alias_returns_none(tmp_path):
 def test_resolve_nonexistent_config_returns_none(tmp_path):
     cfg = tmp_path / "no_such_file"
     assert resolve_instance_id("aws-gpu1", config_path=cfg) is None
+
+
+# ---------------------------------------------------------------------------
+# Cleanup: find_stale_ssh_hosts / cleanup_stale_ssh_hosts
+# ---------------------------------------------------------------------------
+
+
+def test_find_stale_ssh_hosts_finds_orphans(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    stale = find_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg)
+    assert stale == [("i-222bbbb2", "aws-gpu2")]
+
+
+def test_find_stale_ssh_hosts_none_stale(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    stale = find_stale_ssh_hosts({"i-111aaaa1", "i-222bbbb2"}, config_path=cfg)
+    assert stale == []
+
+
+def test_find_stale_ssh_hosts_all_stale(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    stale = find_stale_ssh_hosts(set(), config_path=cfg)
+    assert len(stale) == 2
+    assert ("i-111aaaa1", "aws-gpu1") in stale
+    assert ("i-222bbbb2", "aws-gpu2") in stale
+
+
+def test_find_stale_ssh_hosts_empty_config(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("")
+    stale = find_stale_ssh_hosts(set(), config_path=cfg)
+    assert stale == []
+
+
+def test_cleanup_stale_ssh_hosts_removes(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    results = cleanup_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg)
+    assert len(results) == 1
+    assert results[0].instance_id == "i-222bbbb2"
+    assert results[0].alias == "aws-gpu2"
+    assert results[0].removed is True
+    # Verify it was actually removed from the config
+    content = cfg.read_text()
+    assert "i-222bbbb2" not in content
+    assert "i-111aaaa1" in content
+
+
+def test_cleanup_stale_ssh_hosts_dry_run(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    results = cleanup_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg, dry_run=True)
+    assert len(results) == 1
+    assert results[0].removed is False
+    # Verify config is unchanged
+    content = cfg.read_text()
+    assert "i-222bbbb2" in content
+    assert "i-111aaaa1" in content
+
+
+def test_cleanup_stale_ssh_hosts_no_stale(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    results = cleanup_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg)
+    assert results == []

aws_bootstrap/tests/test_ssh_ebs.py

@@ -0,0 +1,76 @@
+"""Tests for mount_ebs_volume SSH function."""
+
+from __future__ import annotations
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from aws_bootstrap.ssh import mount_ebs_volume
+
+
+KEY_PATH = Path("/home/user/.ssh/id_ed25519.pub")
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_success_format(mock_run):
+    """New volume: SSH command includes mkfs."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    result = mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", format_volume=True)
+
+    assert result is True
+    mock_run.assert_called_once()
+    cmd = mock_run.call_args[0][0]
+    script = cmd[-1]
+    assert "mkfs.ext4" in script
+    assert "/data" in script
+    assert "volabc123" in script  # stripped vol- hyphen
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_success_no_format(mock_run):
+    """Existing volume: SSH command skips mkfs."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    result = mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", format_volume=False)
+
+    assert result is True
+    mock_run.assert_called_once()
+    cmd = mock_run.call_args[0][0]
+    script = cmd[-1]
+    assert "mkfs" not in script
+    assert "/data" in script
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_failure(mock_run):
+    """Non-zero exit code returns False."""
+    mock_run.return_value = MagicMock(returncode=1)
+
+    result = mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123")
+
+    assert result is False
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_custom_port(mock_run):
+    """Non-default port is passed as -p flag."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", port=2222)
+
+    cmd = mock_run.call_args[0][0]
+    assert "-p" in cmd
+    port_idx = cmd.index("-p")
+    assert cmd[port_idx + 1] == "2222"
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_custom_mount_point(mock_run):
+    """Custom mount point appears in the SSH script."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", mount_point="/mnt/data")
+
+    cmd = mock_run.call_args[0][0]
+    script = cmd[-1]
+    assert "/mnt/data" in script