aws-bootstrap-g4dn 0.4.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

aws_bootstrap/tests/test_cli.py

@@ -3,14 +3,14 @@
 from __future__ import annotations
 from datetime import UTC, datetime
 from pathlib import Path
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import botocore.exceptions
 from click.testing import CliRunner
 
 from aws_bootstrap.cli import main
 from aws_bootstrap.gpu import GpuInfo
-from aws_bootstrap.ssh import SSHHostDetails
+from aws_bootstrap.ssh import CleanupResult, SSHHostDetails
 
 
 def test_help():
@@ -170,6 +170,58 @@ def test_terminate_with_confirm(mock_terminate, mock_find, mock_session, mock_re
     assert mock_terminate.call_args[0][1] == ["i-abc123"]
 
 
+@patch("aws_bootstrap.cli.remove_ssh_host", return_value=None)
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.terminate_tagged_instances")
+@patch("aws_bootstrap.cli.resolve_instance_id", return_value="i-abc123")
+def test_terminate_by_alias(mock_resolve, mock_terminate, mock_session, mock_remove_ssh):
+    mock_terminate.return_value = [
+        {
+            "InstanceId": "i-abc123",
+            "PreviousState": {"Name": "running"},
+            "CurrentState": {"Name": "shutting-down"},
+        }
+    ]
+    runner = CliRunner()
+    result = runner.invoke(main, ["terminate", "--yes", "aws-gpu1"])
+    assert result.exit_code == 0
+    assert "Resolved alias 'aws-gpu1' -> i-abc123" in result.output
+    assert "Terminated 1" in result.output
+    mock_resolve.assert_called_once_with("aws-gpu1")
+    mock_terminate.assert_called_once()
+    assert mock_terminate.call_args[0][1] == ["i-abc123"]
+
+
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.resolve_instance_id", return_value=None)
+def test_terminate_unknown_alias_errors(mock_resolve, mock_session):
+    runner = CliRunner()
+    result = runner.invoke(main, ["terminate", "--yes", "aws-gpu99"])
+    assert result.exit_code != 0
+    assert "Could not resolve 'aws-gpu99'" in result.output
+
+
+@patch("aws_bootstrap.cli.remove_ssh_host", return_value=None)
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.terminate_tagged_instances")
+@patch("aws_bootstrap.cli.resolve_instance_id", return_value="i-abc123")
+def test_terminate_by_instance_id_passthrough(mock_resolve, mock_terminate, mock_session, mock_remove_ssh):
+    """Instance IDs are passed through without resolution message."""
+    mock_resolve.return_value = "i-abc123"
+    mock_terminate.return_value = [
+        {
+            "InstanceId": "i-abc123",
+            "PreviousState": {"Name": "running"},
+            "CurrentState": {"Name": "shutting-down"},
+        }
+    ]
+    runner = CliRunner()
+    result = runner.invoke(main, ["terminate", "--yes", "i-abc123"])
+    assert result.exit_code == 0
+    assert "Resolved alias" not in result.output
+    assert "Terminated 1" in result.output
+
+
 @patch("aws_bootstrap.cli.boto3.Session")
 @patch("aws_bootstrap.cli.find_tagged_instances")
 def test_terminate_cancelled(mock_find, mock_session):
@@ -295,7 +347,10 @@ def test_launch_output_shows_ssh_alias(
 ):
     mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
     mock_launch.return_value = {"InstanceId": "i-test123"}
-    mock_wait.return_value = {"PublicIpAddress": "1.2.3.4"}
+    mock_wait.return_value = {
+        "PublicIpAddress": "1.2.3.4",
+        "Placement": {"AvailabilityZone": "us-west-2a"},
+    }
 
     key_path = tmp_path / "id_ed25519.pub"
     key_path.write_text("ssh-ed25519 AAAA test@host")
@@ -747,7 +802,10 @@ def test_launch_python_version_passed_to_setup(
 ):
     mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
     mock_launch.return_value = {"InstanceId": "i-test123"}
-    mock_wait.return_value = {"PublicIpAddress": "1.2.3.4"}
+    mock_wait.return_value = {
+        "PublicIpAddress": "1.2.3.4",
+        "Placement": {"AvailabilityZone": "us-west-2a"},
+    }
 
     key_path = tmp_path / "id_ed25519.pub"
     key_path.write_text("ssh-ed25519 AAAA test@host")
@@ -834,3 +892,365 @@ def test_launch_dry_run_omits_ssh_port_when_default(mock_sg, mock_import, mock_a
     result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--dry-run"])
     assert result.exit_code == 0
     assert "SSH port" not in result.output
+
+
+# ---------------------------------------------------------------------------
+# EBS data volume tests
+# ---------------------------------------------------------------------------
+
+
+def test_launch_help_shows_ebs_options():
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--help"])
+    assert result.exit_code == 0
+    assert "--ebs-storage" in result.output
+    assert "--ebs-volume-id" in result.output
+
+
+def test_launch_ebs_mutual_exclusivity(tmp_path):
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(
+        main, ["launch", "--key-path", str(key_path), "--ebs-storage", "96", "--ebs-volume-id", "vol-abc"]
+    )
+    assert result.exit_code != 0
+    assert "mutually exclusive" in result.output
+
+
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+def test_launch_dry_run_with_ebs_storage(mock_sg, mock_import, mock_ami, mock_session, tmp_path):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--dry-run", "--ebs-storage", "96"])
+    assert result.exit_code == 0
+    assert "96 GB gp3" in result.output
+    assert "/data" in result.output
+
+
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+def test_launch_dry_run_with_ebs_volume_id(mock_sg, mock_import, mock_ami, mock_session, tmp_path):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--dry-run", "--ebs-volume-id", "vol-abc"])
+    assert result.exit_code == 0
+    assert "vol-abc" in result.output
+    assert "/data" in result.output
+
+
+@patch("aws_bootstrap.cli.mount_ebs_volume", return_value=True)
+@patch("aws_bootstrap.cli.attach_ebs_volume")
+@patch("aws_bootstrap.cli.create_ebs_volume", return_value="vol-new123")
+@patch("aws_bootstrap.cli.add_ssh_host", return_value="aws-gpu1")
+@patch("aws_bootstrap.cli.run_remote_setup", return_value=True)
+@patch("aws_bootstrap.cli.wait_for_ssh", return_value=True)
+@patch("aws_bootstrap.cli.wait_instance_ready")
+@patch("aws_bootstrap.cli.launch_instance")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.boto3.Session")
+def test_launch_with_ebs_storage_full_flow(
+    mock_session,
+    mock_ami,
+    mock_import,
+    mock_sg,
+    mock_launch,
+    mock_wait,
+    mock_ssh,
+    mock_setup,
+    mock_add_ssh,
+    mock_create_ebs,
+    mock_attach_ebs,
+    mock_mount_ebs,
+    tmp_path,
+):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+    mock_launch.return_value = {"InstanceId": "i-test123"}
+    mock_wait.return_value = {
+        "PublicIpAddress": "1.2.3.4",
+        "Placement": {"AvailabilityZone": "us-west-2a"},
+    }
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--ebs-storage", "96", "--no-setup"])
+    assert result.exit_code == 0
+    assert "vol-new123" in result.output
+    mock_create_ebs.assert_called_once()
+    mock_attach_ebs.assert_called_once()
+    mock_mount_ebs.assert_called_once()
+    # Verify format_volume=True for new volumes
+    assert mock_mount_ebs.call_args[1]["format_volume"] is True
+
+
+@patch("aws_bootstrap.cli.mount_ebs_volume", return_value=True)
+@patch("aws_bootstrap.cli.attach_ebs_volume")
+@patch("aws_bootstrap.cli.validate_ebs_volume")
+@patch("aws_bootstrap.cli.add_ssh_host", return_value="aws-gpu1")
+@patch("aws_bootstrap.cli.run_remote_setup", return_value=True)
+@patch("aws_bootstrap.cli.wait_for_ssh", return_value=True)
+@patch("aws_bootstrap.cli.wait_instance_ready")
+@patch("aws_bootstrap.cli.launch_instance")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.boto3.Session")
+def test_launch_with_ebs_volume_id_full_flow(
+    mock_session,
+    mock_ami,
+    mock_import,
+    mock_sg,
+    mock_launch,
+    mock_wait,
+    mock_ssh,
+    mock_setup,
+    mock_add_ssh,
+    mock_validate,
+    mock_attach_ebs,
+    mock_mount_ebs,
+    tmp_path,
+):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+    mock_launch.return_value = {"InstanceId": "i-test123"}
+    mock_wait.return_value = {
+        "PublicIpAddress": "1.2.3.4",
+        "Placement": {"AvailabilityZone": "us-west-2a"},
+    }
+    mock_validate.return_value = {"VolumeId": "vol-existing", "Size": 200}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(
+        main, ["launch", "--key-path", str(key_path), "--ebs-volume-id", "vol-existing", "--no-setup"]
+    )
+    assert result.exit_code == 0
+    mock_validate.assert_called_once()
+    mock_attach_ebs.assert_called_once()
+    mock_mount_ebs.assert_called_once()
+    # Verify format_volume=False for existing volumes
+    assert mock_mount_ebs.call_args[1]["format_volume"] is False
+
+
+@patch("aws_bootstrap.cli.mount_ebs_volume", return_value=False)
+@patch("aws_bootstrap.cli.attach_ebs_volume")
+@patch("aws_bootstrap.cli.create_ebs_volume", return_value="vol-new123")
+@patch("aws_bootstrap.cli.add_ssh_host", return_value="aws-gpu1")
+@patch("aws_bootstrap.cli.wait_for_ssh", return_value=True)
+@patch("aws_bootstrap.cli.wait_instance_ready")
+@patch("aws_bootstrap.cli.launch_instance")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.boto3.Session")
+def test_launch_ebs_mount_failure_warns(
+    mock_session,
+    mock_ami,
+    mock_import,
+    mock_sg,
+    mock_launch,
+    mock_wait,
+    mock_ssh,
+    mock_add_ssh,
+    mock_create_ebs,
+    mock_attach_ebs,
+    mock_mount_ebs,
+    tmp_path,
+):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+    mock_launch.return_value = {"InstanceId": "i-test123"}
+    mock_wait.return_value = {
+        "PublicIpAddress": "1.2.3.4",
+        "Placement": {"AvailabilityZone": "us-west-2a"},
+    }
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--ebs-storage", "96", "--no-setup"])
+    # Should succeed despite mount failure (just a warning)
+    assert result.exit_code == 0
+    assert "WARNING" in result.output or "Failed to mount" in result.output
+
+
+def test_terminate_help_shows_keep_ebs():
+    runner = CliRunner()
+    result = runner.invoke(main, ["terminate", "--help"])
+    assert result.exit_code == 0
+    assert "--keep-ebs" in result.output
+
+
+@patch("aws_bootstrap.cli.delete_ebs_volume")
+@patch("aws_bootstrap.cli.find_ebs_volumes_for_instance")
+@patch("aws_bootstrap.cli.remove_ssh_host", return_value=None)
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.find_tagged_instances")
+@patch("aws_bootstrap.cli.terminate_tagged_instances")
+def test_terminate_deletes_ebs_by_default(
+    mock_terminate, mock_find, mock_session, mock_remove_ssh, mock_find_ebs, mock_delete_ebs
+):
+    mock_find.return_value = [
+        {
+            "InstanceId": "i-abc123",
+            "Name": "test",
+            "State": "running",
+            "InstanceType": "g4dn.xlarge",
+            "PublicIp": "1.2.3.4",
+            "LaunchTime": datetime(2025, 1, 1, tzinfo=UTC),
+        }
+    ]
+    mock_terminate.return_value = [
+        {
+            "InstanceId": "i-abc123",
+            "PreviousState": {"Name": "running"},
+            "CurrentState": {"Name": "shutting-down"},
+        }
+    ]
+    mock_find_ebs.return_value = [{"VolumeId": "vol-data1", "Size": 96, "Device": "/dev/sdf", "State": "in-use"}]
+
+    # Mock the ec2 client's get_waiter for volume_available
+    mock_ec2 = mock_session.return_value.client.return_value
+    mock_waiter = MagicMock()
+    mock_ec2.get_waiter.return_value = mock_waiter
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["terminate", "--yes"])
+    assert result.exit_code == 0
+    mock_delete_ebs.assert_called_once_with(mock_ec2, "vol-data1")
+
+
+@patch("aws_bootstrap.cli.delete_ebs_volume")
+@patch("aws_bootstrap.cli.find_ebs_volumes_for_instance")
+@patch("aws_bootstrap.cli.remove_ssh_host", return_value=None)
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.find_tagged_instances")
+@patch("aws_bootstrap.cli.terminate_tagged_instances")
+def test_terminate_keep_ebs_preserves(
+    mock_terminate, mock_find, mock_session, mock_remove_ssh, mock_find_ebs, mock_delete_ebs
+):
+    mock_find.return_value = [
+        {
+            "InstanceId": "i-abc123",
+            "Name": "test",
+            "State": "running",
+            "InstanceType": "g4dn.xlarge",
+            "PublicIp": "1.2.3.4",
+            "LaunchTime": datetime(2025, 1, 1, tzinfo=UTC),
+        }
+    ]
+    mock_terminate.return_value = [
+        {
+            "InstanceId": "i-abc123",
+            "PreviousState": {"Name": "running"},
+            "CurrentState": {"Name": "shutting-down"},
+        }
+    ]
+    mock_find_ebs.return_value = [{"VolumeId": "vol-data1", "Size": 96, "Device": "/dev/sdf", "State": "in-use"}]
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["terminate", "--yes", "--keep-ebs"])
+    assert result.exit_code == 0
+    assert "Preserving EBS volume: vol-data1" in result.output
+    assert "aws-bootstrap launch --ebs-volume-id vol-data1" in result.output
+    mock_delete_ebs.assert_not_called()
+
+
+@patch("aws_bootstrap.cli.find_ebs_volumes_for_instance")
+@patch("aws_bootstrap.cli.get_ssh_host_details", return_value=None)
+@patch("aws_bootstrap.cli.list_ssh_hosts", return_value={})
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_spot_price")
+@patch("aws_bootstrap.cli.find_tagged_instances")
+def test_status_shows_ebs_volumes(mock_find, mock_spot_price, mock_session, mock_ssh_hosts, mock_details, mock_ebs):
+    mock_find.return_value = [
+        {
+            "InstanceId": "i-abc123",
+            "Name": "aws-bootstrap-g4dn.xlarge",
+            "State": "running",
+            "InstanceType": "g4dn.xlarge",
+            "PublicIp": "1.2.3.4",
+            "LaunchTime": datetime(2025, 1, 1, tzinfo=UTC),
+            "Lifecycle": "spot",
+            "AvailabilityZone": "us-west-2a",
+        }
+    ]
+    mock_spot_price.return_value = 0.15
+    mock_ebs.return_value = [{"VolumeId": "vol-data1", "Size": 96, "Device": "/dev/sdf", "State": "in-use"}]
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["status"])
+    assert result.exit_code == 0
+    assert "vol-data1" in result.output
+    assert "96 GB" in result.output
+    assert "/data" in result.output
+
+
+# ---------------------------------------------------------------------------
+# cleanup subcommand
+# ---------------------------------------------------------------------------
+
+
+def test_cleanup_help():
+    runner = CliRunner()
+    result = runner.invoke(main, ["cleanup", "--help"])
+    assert result.exit_code == 0
+    assert "--dry-run" in result.output
+    assert "--yes" in result.output
+    assert "--region" in result.output
+    assert "--profile" in result.output
+
+
+@patch("aws_bootstrap.cli.find_stale_ssh_hosts", return_value=[])
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.find_tagged_instances", return_value=[])
+def test_cleanup_no_stale(mock_find, mock_session, mock_stale):
+    runner = CliRunner()
+    result = runner.invoke(main, ["cleanup"])
+    assert result.exit_code == 0
+    assert "No stale" in result.output
+
+
+@patch("aws_bootstrap.cli.find_stale_ssh_hosts", return_value=[("i-dead1234", "aws-gpu1")])
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.find_tagged_instances", return_value=[])
+def test_cleanup_dry_run(mock_find, mock_session, mock_stale):
+    runner = CliRunner()
+    result = runner.invoke(main, ["cleanup", "--dry-run"])
+    assert result.exit_code == 0
+    assert "Would remove" in result.output
+    assert "aws-gpu1" in result.output
+    assert "i-dead1234" in result.output
+
+
+@patch("aws_bootstrap.cli.cleanup_stale_ssh_hosts")
+@patch("aws_bootstrap.cli.find_stale_ssh_hosts", return_value=[("i-dead1234", "aws-gpu1")])
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.find_tagged_instances", return_value=[])
+def test_cleanup_with_yes(mock_find, mock_session, mock_stale, mock_cleanup):
+    mock_cleanup.return_value = [CleanupResult(instance_id="i-dead1234", alias="aws-gpu1", removed=True)]
+    runner = CliRunner()
+    result = runner.invoke(main, ["cleanup", "--yes"])
+    assert result.exit_code == 0
+    assert "Removed aws-gpu1" in result.output
+    assert "Cleaned up 1" in result.output
+    mock_cleanup.assert_called_once()

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