Add type for build log message

Author: buurro

src/fastapi_cloud_cli/commands/deploy.py

@@ -19,7 +19,7 @@
 from typing_extensions import Annotated
 
 from fastapi_cloud_cli.commands.login import login
-from fastapi_cloud_cli.utils.api import APIClient, BuildLogError, BuildLogType
+from fastapi_cloud_cli.utils.api import APIClient, BuildLogError
 from fastapi_cloud_cli.utils.apps import AppConfig, get_app_config, write_app_config
 from fastapi_cloud_cli.utils.auth import is_logged_in
 from fastapi_cloud_cli.utils.cli import get_rich_toolkit, handle_http_errors
@@ -343,10 +343,10 @@ def _wait_for_deployment(
             for log in client.stream_build_logs(deployment.id):
                 time_elapsed = time.monotonic() - started_at
 
-                if log.type == BuildLogType.message and log.message:
+                if log.type == "message":
                     progress.log(Text.from_ansi(log.message.rstrip()))
 
-                if log.type == BuildLogType.complete:
+                if log.type == "complete":
                     progress.log("")
                     progress.log(
                         f"🐔 Ready the chicken! Your app is ready at [link={deployment.url}]{deployment.url}[/link]"
@@ -360,7 +360,7 @@ def _wait_for_deployment(
 
                     break
 
-                if log.type == BuildLogType.failed:
+                if log.type == "failed":
                     progress.log("")
                     progress.log(
                         f"😔 Oh no! Something went wrong. Check out the logs at [link={deployment.dashboard_url}]{deployment.dashboard_url}[/link]"

src/fastapi_cloud_cli/utils/api.py

@@ -2,11 +2,11 @@
 import time
 from contextlib import contextmanager
 from datetime import timedelta
-from enum import Enum
-from typing import ContextManager, Generator, Optional
+from typing import ContextManager, Generator, Literal, Optional, Union
 
 import httpx
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, Field, TypeAdapter, ValidationError
+from typing_extensions import Annotated
 
 from fastapi_cloud_cli import __version__
 from fastapi_cloud_cli.config import Settings
@@ -18,23 +18,27 @@
 BUILD_LOG_TIMEOUT = timedelta(minutes=5)
 
 
-class BuildLogError(Exception): ...
+class BuildLogError(Exception):
+    pass
 
 
-class BuildLogType(str, Enum):
-    message = "message"
-    complete = "complete"
-    failed = "failed"
-    timeout = "timeout"  # Request closed, reconnect to continue
-    heartbeat = "heartbeat"  # Keepalive signal when no new logs
+class BuildLogLineGeneric(BaseModel):
+    type: Literal["complete", "failed", "timeout", "heartbeat"]
+    id: Optional[str] = None
 
 
-class BuildLogLine(BaseModel):
-    type: BuildLogType
-    message: Optional[str] = None
+class BuildLogLineMessage(BaseModel):
+    type: Literal["message"] = "message"
+    message: str
     id: Optional[str] = None
 
 
+BuildLogLine = Union[BuildLogLineMessage, BuildLogLineGeneric]
+BuildLogAdapter = TypeAdapter[BuildLogLine](
+    Annotated[BuildLogLine, Field(discriminator="type")]
+)
+
+
 @contextmanager
 def attempt(attempt_number: int) -> Generator[None, None, None]:
     def _backoff() -> None:
@@ -132,18 +136,15 @@ def stream_build_logs(
                                 if log_line.id:
                                     last_id = log_line.id
 
-                                if log_line.type == BuildLogType.message:
+                                if log_line.type == "message":
                                     yield log_line
 
-                                if log_line.type in (
-                                    BuildLogType.complete,
-                                    BuildLogType.failed,
-                                ):
+                                if log_line.type in ("complete", "failed"):
                                     yield log_line
 
                                     return
 
-                                if log_line.type == BuildLogType.timeout:
+                                if log_line.type == "timeout":
                                     logger.debug("Received timeout; reconnecting")
                                     break  # Breaks for loop to reconnect
 
@@ -160,7 +161,7 @@ def stream_build_logs(
 
     def _parse_log_line(self, line: str) -> Optional[BuildLogLine]:
         try:
-            return BuildLogLine.model_validate_json(line)
+            return BuildLogAdapter.validate_json(line)
         except ValidationError as e:
             logger.debug("Skipping malformed log: %s (error: %s)", line[:100], e)
             return None

tests/test_api_client.py

@@ -12,7 +12,7 @@
     BUILD_LOG_MAX_RETRIES,
     APIClient,
     BuildLogError,
-    BuildLogType,
+    BuildLogLineMessage,
 )
 from tests.utils import build_logs_response
 
@@ -59,13 +59,13 @@ def test_stream_build_logs_successful(
 
     assert len(logs) == 3
 
-    assert logs[0].type == BuildLogType.message
+    assert logs[0].type == "message"
     assert logs[0].message == "Building..."
 
-    assert logs[1].type == BuildLogType.message
+    assert logs[1].type == "message"
     assert logs[1].message == "Done!"
 
-    assert logs[2].type == BuildLogType.complete
+    assert logs[2].type == "complete"
 
 
 @api_mock
@@ -85,16 +85,16 @@ def test_stream_build_logs_failed(
     logs = list(client.stream_build_logs(deployment_id))
 
     assert len(logs) == 2
-    assert logs[0].type == BuildLogType.message
-    assert logs[1].type == BuildLogType.failed
+    assert logs[0].type == "message"
+    assert logs[1].type == "failed"
 
 
-@pytest.mark.parametrize("terminal_type", [BuildLogType.complete, BuildLogType.failed])
+@pytest.mark.parametrize("terminal_type", ["complete", "failed"])
 @api_mock
 def test_stream_build_logs_stop_after_terminal_state(
     logs_route: respx.Route,
     client: APIClient,
-    terminal_type: BuildLogType,
+    terminal_type: str,
     deployment_id: str,
 ) -> None:
     logs_route.mock(
@@ -111,7 +111,7 @@ def test_stream_build_logs_stop_after_terminal_state(
     logs = list(client.stream_build_logs(deployment_id))
 
     assert len(logs) == 2
-    assert logs[0].type == BuildLogType.message
+    assert logs[0].type == "message"
     assert logs[1].type == terminal_type
 
 
@@ -125,7 +125,7 @@ def test_stream_build_logs_internal_messages_are_skipped(
         return_value=Response(
             200,
             content=build_logs_response(
-                {"type": BuildLogType.heartbeat, "id": "1"},
+                {"type": "heartbeat", "id": "1"},
                 {"type": "message", "message": "Continuing...", "id": "2"},
                 {"type": "complete", "id": "3"},
             ),
@@ -135,8 +135,8 @@ def test_stream_build_logs_internal_messages_are_skipped(
     logs = list(client.stream_build_logs(deployment_id))
 
     assert len(logs) == 2
-    assert logs[0].type == BuildLogType.message
-    assert logs[1].type == BuildLogType.complete
+    assert logs[0].type == "message"
+    assert logs[1].type == "complete"
 
 
 @api_mock
@@ -156,8 +156,8 @@ def test_stream_build_logs_malformed_json_is_skipped(
     logs = list(client.stream_build_logs(deployment_id))
 
     assert len(logs) == 2
-    assert logs[0].type == BuildLogType.message
-    assert logs[1].type == BuildLogType.complete
+    assert logs[0].type == "message"
+    assert logs[1].type == "complete"
 
 
 @api_mock
@@ -179,8 +179,8 @@ def test_stream_build_logs_unknown_log_type_is_skipped(
 
     # Unknown type should be filtered out
     assert len(logs) == 2
-    assert logs[0].type == BuildLogType.message
-    assert logs[1].type == BuildLogType.complete
+    assert logs[0].type == "message"
+    assert logs[1].type == "complete"
 
 
 @pytest.mark.parametrize(
@@ -211,6 +211,7 @@ def test_stream_build_logs_network_error_retry(
         logs = list(client.stream_build_logs(deployment_id))
 
     assert len(logs) == 2
+    assert logs[0].type == "message"
     assert logs[0].message == "Success after retry"
 
 
@@ -232,7 +233,7 @@ def test_stream_build_logs_server_error_retry(
         logs = list(client.stream_build_logs(deployment_id))
 
     assert len(logs) == 1
-    assert logs[0].type == BuildLogType.complete
+    assert logs[0].type == "complete"
 
 
 @api_mock
@@ -277,8 +278,8 @@ def test_stream_build_logs_empty_lines_are_skipped(
     logs = list(client.stream_build_logs(deployment_id))
 
     assert len(logs) == 2
-    assert logs[0].type == BuildLogType.message
-    assert logs[1].type == BuildLogType.complete
+    assert logs[0].type == "message"
+    assert logs[1].type == "complete"
 
 
 @respx.mock(base_url=settings.base_api_url)
@@ -318,11 +319,11 @@ def test_stream_build_logs_continue_after_timeout(
     logs = client.stream_build_logs(deployment_id)
 
     with patch("time.sleep"):
-        assert next(logs).message == "message 1"
-        assert next(logs).message == "message 2"
-        assert next(logs).message == "message 3"
-        assert next(logs).message == "message 4"
-        assert next(logs).type == BuildLogType.complete
+        assert next(logs) == BuildLogLineMessage(message="message 1", id="1")
+        assert next(logs) == BuildLogLineMessage(message="message 2", id="2")
+        assert next(logs) == BuildLogLineMessage(message="message 3", id="3")
+        assert next(logs) == BuildLogLineMessage(message="message 4", id="4")
+        assert next(logs).type == "complete"
 
 
 @api_mock