breakpilot-pwa/backend/tests/test_infra/test_vast_client.py

"""
Tests fuer VastAIClient.

Testet den vast.ai REST API Client.
"""

import pytest
from unittest.mock import AsyncMock, patch, MagicMock
from datetime import datetime, timezone

from infra.vast_client import (
    VastAIClient,
    InstanceInfo,
    InstanceStatus,
    AccountInfo,
)


class TestInstanceStatus:
    """Tests fuer InstanceStatus Enum."""

    def test_status_values(self):
        """Test alle Status-Werte."""
        assert InstanceStatus.RUNNING.value == "running"
        assert InstanceStatus.STOPPED.value == "stopped"
        assert InstanceStatus.EXITED.value == "exited"
        assert InstanceStatus.LOADING.value == "loading"
        assert InstanceStatus.SCHEDULING.value == "scheduling"
        assert InstanceStatus.CREATING.value == "creating"
        assert InstanceStatus.UNKNOWN.value == "unknown"


class TestInstanceInfo:
    """Tests fuer InstanceInfo Dataclass."""

    def test_create_basic(self):
        """Test grundlegende Erstellung."""
        info = InstanceInfo(
            id=12345,
            status=InstanceStatus.RUNNING,
        )
        assert info.id == 12345
        assert info.status == InstanceStatus.RUNNING
        assert info.gpu_name is None
        assert info.num_gpus == 1

    def test_from_api_response_running(self):
        """Test Parse von API Response (running)."""
        api_data = {
            "id": 67890,
            "actual_status": "running",
            "machine_id": 111,
            "gpu_name": "RTX 3090",
            "num_gpus": 1,
            "gpu_ram": 24.0,
            "cpu_ram": 64.0,
            "disk_space": 200.0,
            "dph_total": 0.45,
            "public_ipaddr": "192.168.1.100",
            "ports": {
                "8001/tcp": [{"HostIp": "0.0.0.0", "HostPort": "12345"}],
            },
            "label": "llm-server",
            "start_date": 1702900800,  # Unix timestamp
        }

        info = InstanceInfo.from_api_response(api_data)

        assert info.id == 67890
        assert info.status == InstanceStatus.RUNNING
        assert info.gpu_name == "RTX 3090"
        assert info.dph_total == 0.45
        assert info.public_ipaddr == "192.168.1.100"
        assert info.label == "llm-server"

    def test_from_api_response_stopped(self):
        """Test Parse von API Response (gestoppt)."""
        api_data = {
            "id": 11111,
            "actual_status": "exited",
        }

        info = InstanceInfo.from_api_response(api_data)

        assert info.status == InstanceStatus.EXITED

    def test_from_api_response_unknown_status(self):
        """Test Parse von unbekanntem Status."""
        api_data = {
            "id": 22222,
            "actual_status": "weird_status",
        }

        info = InstanceInfo.from_api_response(api_data)

        assert info.status == InstanceStatus.UNKNOWN

    def test_get_endpoint_url_with_port_mapping(self):
        """Test Endpoint URL mit Port-Mapping."""
        info = InstanceInfo(
            id=12345,
            status=InstanceStatus.RUNNING,
            public_ipaddr="10.0.0.1",
            ports={
                "8001/tcp": [{"HostIp": "0.0.0.0", "HostPort": "54321"}],
            },
        )

        url = info.get_endpoint_url(8001)
        assert url == "http://10.0.0.1:54321"

    def test_get_endpoint_url_fallback(self):
        """Test Endpoint URL Fallback ohne Port-Mapping."""
        info = InstanceInfo(
            id=12345,
            status=InstanceStatus.RUNNING,
            public_ipaddr="10.0.0.2",
            ports={},
        )

        url = info.get_endpoint_url(8001)
        assert url == "http://10.0.0.2:8001"

    def test_get_endpoint_url_no_ip(self):
        """Test Endpoint URL ohne IP."""
        info = InstanceInfo(
            id=12345,
            status=InstanceStatus.RUNNING,
            public_ipaddr=None,
        )

        url = info.get_endpoint_url(8001)
        assert url is None

    def test_to_dict(self):
        """Test Serialisierung."""
        info = InstanceInfo(
            id=12345,
            status=InstanceStatus.RUNNING,
            gpu_name="RTX 4090",
            dph_total=0.75,
        )

        data = info.to_dict()

        assert data["id"] == 12345
        assert data["status"] == "running"
        assert data["gpu_name"] == "RTX 4090"
        assert data["dph_total"] == 0.75


class TestAccountInfo:
    """Tests fuer AccountInfo Dataclass."""

    def test_create_basic(self):
        """Test grundlegende Erstellung."""
        info = AccountInfo(
            credit=25.50,
            balance=0.0,
            total_spend=10.25,
            username="testuser",
            email="test@example.com",
            has_billing=True,
        )
        assert info.credit == 25.50
        assert info.balance == 0.0
        assert info.total_spend == 10.25
        assert info.username == "testuser"
        assert info.email == "test@example.com"
        assert info.has_billing is True

    def test_from_api_response_complete(self):
        """Test Parse von vollstaendiger API Response."""
        api_data = {
            "credit": 23.87674017153,
            "balance": 0.0,
            "total_spend": -1.1732598284700013,  # API gibt negativ zurueck
            "username": "benjamin",
            "email": "benjamin@example.com",
            "has_billing": True,
        }

        info = AccountInfo.from_api_response(api_data)

        assert info.credit == 23.87674017153
        assert info.balance == 0.0
        assert info.total_spend == 1.1732598284700013  # Sollte positiv sein (abs)
        assert info.username == "benjamin"
        assert info.email == "benjamin@example.com"
        assert info.has_billing is True

    def test_from_api_response_minimal(self):
        """Test Parse von minimaler API Response."""
        api_data = {}

        info = AccountInfo.from_api_response(api_data)

        assert info.credit == 0.0
        assert info.balance == 0.0
        assert info.total_spend == 0.0
        assert info.username == ""
        assert info.email == ""
        assert info.has_billing is False

    def test_from_api_response_partial(self):
        """Test Parse von teilweiser API Response."""
        api_data = {
            "credit": 50.0,
            "username": "partial_user",
        }

        info = AccountInfo.from_api_response(api_data)

        assert info.credit == 50.0
        assert info.username == "partial_user"
        assert info.email == ""
        assert info.total_spend == 0.0

    def test_to_dict(self):
        """Test Serialisierung zu Dictionary."""
        info = AccountInfo(
            credit=100.0,
            balance=5.0,
            total_spend=25.0,
            username="dictuser",
            email="dict@test.com",
            has_billing=True,
        )

        data = info.to_dict()

        assert data["credit"] == 100.0
        assert data["balance"] == 5.0
        assert data["total_spend"] == 25.0
        assert data["username"] == "dictuser"
        assert data["email"] == "dict@test.com"
        assert data["has_billing"] is True

    def test_total_spend_negative_to_positive(self):
        """Test dass negative total_spend Werte positiv werden."""
        api_data = {
            "total_spend": -99.99,
        }

        info = AccountInfo.from_api_response(api_data)

        assert info.total_spend == 99.99


class TestVastAIClient:
    """Tests fuer VastAIClient."""

    def test_init(self):
        """Test Client Initialisierung."""
        client = VastAIClient(api_key="test-key", timeout=60.0)

        assert client.api_key == "test-key"
        assert client.timeout == 60.0
        assert client._client is None

    def test_build_url(self):
        """Test URL Building."""
        client = VastAIClient(api_key="my-api-key")

        url = client._build_url("/instances/")
        assert url == "https://console.vast.ai/api/v0/instances/?api_key=my-api-key"

        url2 = client._build_url("/instances/?foo=bar")
        assert url2 == "https://console.vast.ai/api/v0/instances/?foo=bar&api_key=my-api-key"

    @pytest.mark.asyncio
    async def test_list_instances(self):
        """Test Liste aller Instanzen."""
        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "instances": [
                {"id": 1, "actual_status": "running", "gpu_name": "RTX 3090"},
                {"id": 2, "actual_status": "exited"},  # exited statt stopped
            ]
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.get.return_value = mock_response
            mock_get.return_value = mock_client

            instances = await client.list_instances()

        assert len(instances) == 2
        assert instances[0].id == 1
        assert instances[0].status == InstanceStatus.RUNNING
        assert instances[1].id == 2
        assert instances[1].status == InstanceStatus.EXITED

    @pytest.mark.asyncio
    async def test_get_instance(self):
        """Test einzelne Instanz abrufen."""
        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "instances": [
                {"id": 12345, "actual_status": "running", "dph_total": 0.5}
            ]
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.get.return_value = mock_response
            mock_get.return_value = mock_client

            instance = await client.get_instance(12345)

        assert instance is not None
        assert instance.id == 12345
        assert instance.status == InstanceStatus.RUNNING
        assert instance.dph_total == 0.5

    @pytest.mark.asyncio
    async def test_get_instance_not_found(self):
        """Test Instanz nicht gefunden."""
        import httpx

        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 404

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.get.side_effect = httpx.HTTPStatusError(
                "Not Found", request=MagicMock(), response=mock_response
            )
            mock_get.return_value = mock_client

            instance = await client.get_instance(99999)

        assert instance is None

    @pytest.mark.asyncio
    async def test_start_instance(self):
        """Test Instanz starten."""
        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.put.return_value = mock_response
            mock_get.return_value = mock_client

            success = await client.start_instance(12345)

        assert success is True
        mock_client.put.assert_called_once()
        call_args = mock_client.put.call_args
        assert call_args[1]["json"] == {"state": "running"}

    @pytest.mark.asyncio
    async def test_stop_instance(self):
        """Test Instanz stoppen."""
        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.put.return_value = mock_response
            mock_get.return_value = mock_client

            success = await client.stop_instance(12345)

        assert success is True
        call_args = mock_client.put.call_args
        assert call_args[1]["json"] == {"state": "stopped"}

    @pytest.mark.asyncio
    async def test_stop_instance_failure(self):
        """Test Instanz stoppen fehlschlaegt."""
        import httpx

        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 500

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.put.side_effect = httpx.HTTPStatusError(
                "Error", request=MagicMock(), response=mock_response
            )
            mock_get.return_value = mock_client

            success = await client.stop_instance(12345)

        assert success is False

    @pytest.mark.asyncio
    async def test_destroy_instance(self):
        """Test Instanz loeschen."""
        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.delete.return_value = mock_response
            mock_get.return_value = mock_client

            success = await client.destroy_instance(12345)

        assert success is True
        mock_client.delete.assert_called_once()

    @pytest.mark.asyncio
    async def test_set_label(self):
        """Test Label setzen."""
        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.put.return_value = mock_response
            mock_get.return_value = mock_client

            success = await client.set_label(12345, "my-label")

        assert success is True
        call_args = mock_client.put.call_args
        assert call_args[1]["json"] == {"label": "my-label"}

    @pytest.mark.asyncio
    async def test_close(self):
        """Test Client schliessen."""
        client = VastAIClient(api_key="test-key")

        # Erstelle Client
        await client._get_client()
        assert client._client is not None

        # Schliesse
        await client.close()
        assert client._client is None

    @pytest.mark.asyncio
    async def test_get_account_info_success(self):
        """Test Account-Info erfolgreich abrufen."""
        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "credit": 23.87674017153,
            "balance": 0.0,
            "total_spend": -1.1732598284700013,
            "username": "testuser",
            "email": "test@vast.ai",
            "has_billing": True,
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.get.return_value = mock_response
            mock_get.return_value = mock_client

            account = await client.get_account_info()

        assert account is not None
        assert account.credit == 23.87674017153
        assert account.total_spend == 1.1732598284700013  # abs()
        assert account.username == "testuser"
        assert account.email == "test@vast.ai"
        assert account.has_billing is True

    @pytest.mark.asyncio
    async def test_get_account_info_api_error(self):
        """Test Account-Info bei API-Fehler."""
        import httpx

        client = VastAIClient(api_key="test-key")

        mock_response = MagicMock()
        mock_response.status_code = 401

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.get.side_effect = httpx.HTTPStatusError(
                "Unauthorized", request=MagicMock(), response=mock_response
            )
            mock_get.return_value = mock_client

            account = await client.get_account_info()

        assert account is None

    @pytest.mark.asyncio
    async def test_get_account_info_network_error(self):
        """Test Account-Info bei Netzwerk-Fehler."""
        client = VastAIClient(api_key="test-key")

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.get.side_effect = Exception("Network error")
            mock_get.return_value = mock_client

            account = await client.get_account_info()

        assert account is None

    @pytest.mark.asyncio
    async def test_get_account_info_url(self):
        """Test dass get_account_info den korrekten Endpoint aufruft."""
        client = VastAIClient(api_key="my-test-key")

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"credit": 10.0}
        mock_response.raise_for_status = MagicMock()

        with patch.object(client, "_get_client") as mock_get:
            mock_client = AsyncMock()
            mock_client.get.return_value = mock_response
            mock_get.return_value = mock_client

            await client.get_account_info()

        # Pruefe dass /users/current/ aufgerufen wurde
        call_args = mock_client.get.call_args
        called_url = call_args[0][0]
        assert "/users/current/" in called_url
        assert "api_key=my-test-key" in called_url