"""Tests for main module helper functions.""" from pathlib import Path from unittest.mock import MagicMock, patch import pytest import typer from dropkit.main import ( _resize_hibernated_snapshot, add_temporary_ssh_rule, build_droplet_tags, complete_droplet_or_snapshot_name, find_snapshot_action, get_droplet_name_from_snapshot, get_snapshot_name, get_ssh_hostname, get_user_tag, is_droplet_tailscale_locked, prepare_for_hibernate, ) class TestGetSnapshotName: """Tests for get_snapshot_name function.""" def test_simple_name(self): """Test snapshot name for simple droplet name.""" assert get_snapshot_name("myvm") == "dropkit-myvm" def test_name_with_hyphen(self): """Test snapshot name for droplet name with hyphen.""" assert get_snapshot_name("my-droplet") == "dropkit-my-droplet" def test_name_with_numbers(self): """Test snapshot name for droplet name with numbers.""" assert get_snapshot_name("test123") == "dropkit-test123" def test_empty_name(self): """Test snapshot name for empty droplet name.""" assert get_snapshot_name("") == "dropkit-" class TestGetDropletNameFromSnapshot: """Tests for get_droplet_name_from_snapshot function.""" def test_valid_dropkit_snapshot(self): """Test extracting droplet name from valid dropkit snapshot.""" assert get_droplet_name_from_snapshot("dropkit-myvm") == "myvm" def test_snapshot_with_hyphen_in_name(self): """Test extracting droplet name with hyphens.""" assert get_droplet_name_from_snapshot("dropkit-my-droplet") == "my-droplet" def test_non_dropkit_snapshot(self): """Test with non-dropkit snapshot name.""" assert get_droplet_name_from_snapshot("other-snapshot") is None def test_partial_prefix(self): """Test with partial prefix (should not match).""" assert get_droplet_name_from_snapshot("dropkit") is None def test_different_prefix(self): """Test with different prefix.""" assert get_droplet_name_from_snapshot("snapshot-myvm") is None def test_empty_after_prefix(self): """Test snapshot name that is just the prefix.""" assert get_droplet_name_from_snapshot("dropkit-") == "" class TestGetSshHostname: """Tests for get_ssh_hostname function.""" def test_simple_name(self): """Test SSH hostname for simple droplet name.""" assert get_ssh_hostname("myvm") == "dropkit.myvm" def test_name_with_hyphen(self): """Test SSH hostname for droplet name with hyphen.""" assert get_ssh_hostname("my-droplet") == "dropkit.my-droplet" class TestGetUserTag: """Tests for get_user_tag function.""" def test_simple_username(self): """Test user tag for simple username.""" assert get_user_tag("john") == "owner:john" def test_username_with_underscore(self): """Test user tag for username with underscore.""" assert get_user_tag("john_doe") == "owner:john_doe" class TestBuildDropletTags: """Tests for build_droplet_tags function.""" def test_no_extra_tags(self): """Test building tags without extra tags.""" tags = build_droplet_tags("john") assert tags == ["owner:john", "firewall"] def test_with_extra_tags(self): """Test building tags with extra tags.""" tags = build_droplet_tags("john", ["production", "webserver"]) assert tags == ["owner:john", "firewall", "production", "webserver"] def test_extra_tags_no_duplicates(self): """Test that duplicate tags are not added.""" tags = build_droplet_tags("john", ["firewall", "production"]) assert tags == ["owner:john", "firewall", "production"] def test_empty_extra_tags(self): """Test with empty extra tags list.""" tags = build_droplet_tags("john", []) assert tags == ["owner:john", "firewall"] def test_none_extra_tags(self): """Test with None extra tags.""" tags = build_droplet_tags("john", None) assert tags == ["owner:john", "firewall"] class TestFindSnapshotAction: """Tests for find_snapshot_action function.""" def test_finds_snapshot_action(self): """Test finding a snapshot action in the actions list.""" mock_api = MagicMock() mock_api.list_droplet_actions.return_value = [ {"id": 1, "type": "power_off", "status": "completed"}, {"id": 2, "type": "snapshot", "status": "in-progress"}, {"id": 3, "type": "power_on", "status": "completed"}, ] result = find_snapshot_action(mock_api, 12345) assert result is not None assert result["id"] == 2 assert result["type"] == "snapshot" mock_api.list_droplet_actions.assert_called_once_with(12345) def test_returns_first_snapshot_action(self): """Test that it returns the first (most recent) snapshot action.""" mock_api = MagicMock() mock_api.list_droplet_actions.return_value = [ {"id": 10, "type": "snapshot", "status": "in-progress"}, {"id": 5, "type": "snapshot", "status": "completed"}, ] result = find_snapshot_action(mock_api, 12345) assert result is not None assert result["id"] == 10 def test_returns_none_when_no_snapshot_action(self): """Test returning None when no snapshot action exists.""" mock_api = MagicMock() mock_api.list_droplet_actions.return_value = [ {"id": 1, "type": "power_off", "status": "completed"}, {"id": 2, "type": "power_on", "status": "completed"}, ] result = find_snapshot_action(mock_api, 12345) assert result is None def test_returns_none_when_empty_actions(self): """Test returning None when actions list is empty.""" mock_api = MagicMock() mock_api.list_droplet_actions.return_value = [] result = find_snapshot_action(mock_api, 12345) assert result is None @pytest.fixture def temp_ssh_config(tmp_path): """Create a temporary SSH config file.""" ssh_dir = tmp_path / ".ssh" ssh_dir.mkdir(mode=0o700) config_path = ssh_dir / "config" return str(config_path) class TestIsDropletTailscaleLocked: """Tests for is_droplet_tailscale_locked function.""" def test_tailscale_ip_returns_true(self, temp_ssh_config): """Test returns True when SSH config has Tailscale IP.""" # Create SSH config with Tailscale IP Path(temp_ssh_config).write_text("""Host dropkit.myvm HostName 100.80.123.45 User ubuntu """) mock_config = MagicMock() mock_config.ssh.config_path = temp_ssh_config result = is_droplet_tailscale_locked(mock_config, "myvm") assert result is True def test_public_ip_returns_false(self, temp_ssh_config): """Test returns False when SSH config has public IP.""" # Create SSH config with public IP Path(temp_ssh_config).write_text("""Host dropkit.myvm HostName 192.168.1.100 User ubuntu """) mock_config = MagicMock() mock_config.ssh.config_path = temp_ssh_config result = is_droplet_tailscale_locked(mock_config, "myvm") assert result is False def test_missing_entry_returns_false(self, temp_ssh_config): """Test returns False when SSH config has no entry for droplet.""" # Create SSH config without the target host Path(temp_ssh_config).write_text("""Host dropkit.othervm HostName 100.80.123.45 User ubuntu """) mock_config = MagicMock() mock_config.ssh.config_path = temp_ssh_config result = is_droplet_tailscale_locked(mock_config, "myvm") assert result is False def test_nonexistent_config_file_returns_false(self, temp_ssh_config): """Test returns False when SSH config file doesn't exist.""" mock_config = MagicMock() mock_config.ssh.config_path = "/nonexistent/path/config" result = is_droplet_tailscale_locked(mock_config, "myvm") assert result is False class TestAddTemporarySshRule: """Tests for add_temporary_ssh_rule function.""" @patch("dropkit.main.subprocess.run") def test_success(self, mock_run): """Test successful SSH rule addition.""" mock_run.return_value = MagicMock(returncode=0, stdout="", stderr="") result = add_temporary_ssh_rule("dropkit.myvm") assert result is True mock_run.assert_called_once() call_args = mock_run.call_args assert "dropkit.myvm" in call_args[0][0] assert "sudo ufw allow in on eth0 to any port 22" in call_args[0][0] @patch("dropkit.main.subprocess.run") def test_failure(self, mock_run): """Test failed SSH rule addition.""" mock_run.return_value = MagicMock(returncode=1, stdout="", stderr="error") result = add_temporary_ssh_rule("dropkit.myvm") assert result is False @patch("dropkit.main.subprocess.run") def test_timeout(self, mock_run): """Test SSH timeout.""" import subprocess mock_run.side_effect = subprocess.TimeoutExpired("ssh", 30) result = add_temporary_ssh_rule("dropkit.myvm") assert result is False class TestPrepareForHibernate: """Tests for prepare_for_hibernate function.""" def test_not_tailscale_locked_returns_false(self, temp_ssh_config): """Test returns False when droplet is not Tailscale locked.""" # Create SSH config with public IP (not Tailscale locked) Path(temp_ssh_config).write_text("""Host dropkit.myvm HostName 192.168.1.100 User ubuntu """) mock_config = MagicMock() mock_config.ssh.config_path = temp_ssh_config mock_api = MagicMock() mock_droplet = {"networks": {"v4": [{"type": "public", "ip_address": "192.168.1.100"}]}} result = prepare_for_hibernate(mock_config, mock_api, mock_droplet, "myvm") assert result is False @patch("dropkit.main.tailscale_logout") @patch("dropkit.main.add_temporary_ssh_rule") @patch("dropkit.main.add_ssh_host") def test_tailscale_locked_returns_true( self, mock_add_ssh_host, mock_add_temp_rule, mock_logout, temp_ssh_config ): """Test returns True when droplet is Tailscale locked.""" # Create SSH config with Tailscale IP Path(temp_ssh_config).write_text("""Host dropkit.myvm HostName 100.80.123.45 User ubuntu """) mock_config = MagicMock() mock_config.ssh.config_path = temp_ssh_config mock_config.ssh.identity_file = "~/.ssh/id_ed25519" mock_api = MagicMock() mock_api.get_username.return_value = "testuser" mock_droplet = {"networks": {"v4": [{"type": "public", "ip_address": "203.0.113.50"}]}} mock_add_temp_rule.return_value = True mock_logout.return_value = True result = prepare_for_hibernate(mock_config, mock_api, mock_droplet, "myvm") assert result is True mock_add_temp_rule.assert_called_once() mock_logout.assert_called_once() mock_add_ssh_host.assert_called_once() @patch("dropkit.main.tailscale_logout") @patch("dropkit.main.add_temporary_ssh_rule") @patch("dropkit.main.add_ssh_host") def test_temp_rule_failure_skips_logout( self, mock_add_ssh_host, mock_add_temp_rule, mock_logout, temp_ssh_config ): """Test returns True but skips logout if temp rule fails (safety).""" # Create SSH config with Tailscale IP Path(temp_ssh_config).write_text("""Host dropkit.myvm HostName 100.80.123.45 User ubuntu """) mock_config = MagicMock() mock_config.ssh.config_path = temp_ssh_config mock_config.ssh.identity_file = "~/.ssh/id_ed25519" mock_api = MagicMock() mock_api.get_username.return_value = "testuser" mock_droplet = {"networks": {"v4": [{"type": "public", "ip_address": "203.0.113.50"}]}} mock_add_temp_rule.return_value = False # Simulating failure result = prepare_for_hibernate(mock_config, mock_api, mock_droplet, "myvm") # Should still return True because we detected Tailscale lockdown assert result is True # But logout should NOT be called (safety - need public IP fallback first) mock_logout.assert_not_called() # And SSH config should NOT be updated (early return) mock_add_ssh_host.assert_not_called() class TestResizeHibernatedSnapshot: """Tests for _resize_hibernated_snapshot function.""" def test_no_snapshot_id_exits(self): """Test exits with error when snapshot has no ID.""" mock_api = MagicMock() snapshot = {"tags": ["size:s-1vcpu-1gb"]} with pytest.raises(typer.Exit): _resize_hibernated_snapshot(mock_api, snapshot, "myvm", "s-2vcpu-4gb") def test_no_size_tag_exits(self): """Test exits with error when snapshot has no size: tag.""" mock_api = MagicMock() snapshot = {"id": "12345", "tags": ["owner:testuser", "firewall"]} with pytest.raises(typer.Exit): _resize_hibernated_snapshot(mock_api, snapshot, "myvm", "s-2vcpu-4gb") @patch("dropkit.main.Prompt.ask", return_value="yes") def test_same_size_exits(self, mock_prompt): """Test exits when new size matches current size.""" mock_api = MagicMock() snapshot = {"id": "12345", "tags": ["size:s-1vcpu-1gb", "owner:testuser"]} with pytest.raises(typer.Exit): _resize_hibernated_snapshot(mock_api, snapshot, "myvm", "s-1vcpu-1gb") @patch("dropkit.main.Prompt.ask", return_value="yes") def test_successful_resize_swaps_tags(self, mock_prompt): """Test successful resize creates new tag, tags resource, then untags old.""" mock_api = MagicMock() mock_api.get_available_sizes.return_value = [ {"slug": "s-2vcpu-4gb", "vcpus": 2, "memory": 4096, "disk": 80, "price_monthly": 24}, ] snapshot = {"id": "12345", "tags": ["size:s-1vcpu-1gb", "owner:testuser"]} _resize_hibernated_snapshot(mock_api, snapshot, "myvm", "s-2vcpu-4gb") # Verify tag operations: add new first, then remove old mock_api.create_tag.assert_called_once_with("size:s-2vcpu-4gb") mock_api.tag_resource.assert_called_once_with("size:s-2vcpu-4gb", "12345", "image") mock_api.untag_resource.assert_called_once_with("size:s-1vcpu-1gb", "12345", "image") @patch("dropkit.main.Prompt.ask", return_value="no") def test_cancelled_resize_no_api_calls(self, mock_prompt): """Test that cancelling resize makes no tag API calls.""" mock_api = MagicMock() mock_api.get_available_sizes.return_value = [ {"slug": "s-2vcpu-4gb", "vcpus": 2, "memory": 4096, "disk": 80, "price_monthly": 24}, ] snapshot = {"id": "12345", "tags": ["size:s-1vcpu-1gb", "owner:testuser"]} with pytest.raises(typer.Exit): _resize_hibernated_snapshot(mock_api, snapshot, "myvm", "s-2vcpu-4gb") mock_api.create_tag.assert_not_called() mock_api.tag_resource.assert_not_called() mock_api.untag_resource.assert_not_called() @patch("dropkit.main.Prompt.ask", return_value="yes") def test_invalid_size_exits(self, mock_prompt): """Test exits when provided size slug doesn't exist.""" mock_api = MagicMock() mock_api.get_available_sizes.return_value = [ {"slug": "s-1vcpu-1gb", "vcpus": 1, "memory": 1024, "disk": 25, "price_monthly": 6}, ] snapshot = {"id": "12345", "tags": ["size:s-1vcpu-1gb", "owner:testuser"]} with pytest.raises(typer.Exit): _resize_hibernated_snapshot(mock_api, snapshot, "myvm", "nonexistent-size") class TestCompleteDropletOrSnapshotName: """Tests for complete_droplet_or_snapshot_name function.""" @patch("dropkit.main.complete_snapshot_name", return_value=["snap-vm"]) @patch("dropkit.main.complete_droplet_name", return_value=["live-vm"]) def test_combines_both_sources(self, mock_droplet, mock_snapshot): """Test that results from both completers are combined.""" result = complete_droplet_or_snapshot_name("") assert "live-vm" in result assert "snap-vm" in result @patch("dropkit.main.complete_snapshot_name", return_value=["shared-vm"]) @patch("dropkit.main.complete_droplet_name", return_value=["shared-vm"]) def test_deduplicates(self, mock_droplet, mock_snapshot): """Test that duplicate names appear only once.""" result = complete_droplet_or_snapshot_name("") assert result.count("shared-vm") == 1 @patch("dropkit.main.complete_snapshot_name", return_value=["snap-vm"]) @patch("dropkit.main.complete_droplet_name", return_value=["live-vm"]) def test_droplets_first(self, mock_droplet, mock_snapshot): """Test that live droplet names appear before snapshot names.""" result = complete_droplet_or_snapshot_name("") assert result.index("live-vm") < result.index("snap-vm")