core: services: disk_usage: Add dynamic disk speed test

Signed-off-by: Patrick José Pereira <patrickelectric@gmail.com>

Author: patrickelectric

core/services/disk_usage/main.py

@@ -9,12 +9,14 @@
 import time
 from functools import wraps
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from commonwealth.utils.apis import GenericErrorHandlingRoute, PrettyJSONResponse
 from commonwealth.utils.logs import InterceptHandler, init_logger
 from commonwealth.utils.sentry_config import init_sentry_async
+from commonwealth.utils.streaming import streamer
 from fastapi import APIRouter, FastAPI, HTTPException, Query, status
+from fastapi.responses import StreamingResponse
 from fastapi_versioning import VersionedFastAPI, versioned_api_route
 from loguru import logger
 from pydantic import BaseModel, Field
@@ -56,6 +58,12 @@ class DiskSpeedResult(BaseModel):
     error: Optional[str] = Field(None, description="Error message if test failed")
 
 
+class DiskSpeedTestPoint(BaseModel):
+    size_mb: int = Field(..., description="Test size in MB")
+    write_speed: Optional[float] = Field(None, description="Write speed in MiB/s")
+    read_speed: Optional[float] = Field(None, description="Read speed in MiB/s")
+
+
 DiskNode.update_forward_refs()
 
 
@@ -303,20 +311,8 @@ def parse_disktest_speed(output: str) -> tuple[Optional[float], Optional[float],
     return write_speed, read_speed, seed
 
 
-# pylint: disable=too-many-locals
-@disk_router.get(
-    "/speed",
-    response_model=DiskSpeedResult,
-    summary="Run disk speed test using disktest binary.",
-)
-@to_http_exception
-async def disk_speed(
-    size_bytes: int = Query(
-        1024 * 1024 * 1024,
-        ge=1024 * 1024,
-        description="Number of bytes to test (default 1 GiB).",
-    ),
-) -> DiskSpeedResult:
+async def run_single_speed_test(size_bytes: int) -> DiskSpeedResult:
+    """Run a single disk speed test and return the result."""
     disktest_binary = "disktest"
     temp_file_path: Optional[Path] = None
 
@@ -408,6 +404,58 @@ async def disk_speed(
                 logger.warning(f"Failed to clean up temporary file {temp_file_path}: {e}")
 
 
+# pylint: disable=too-many-locals
+@disk_router.get(
+    "/speed",
+    response_model=DiskSpeedResult,
+    summary="Run disk speed test using disktest binary.",
+)
+@to_http_exception
+async def disk_speed(
+    size_bytes: int = Query(
+        1024 * 1024 * 1024,
+        ge=1024 * 1024,
+        description="Number of bytes to test (default 1 GiB).",
+    ),
+) -> DiskSpeedResult:
+    return await run_single_speed_test(size_bytes)
+
+
+async def multi_size_speed_test_generator() -> AsyncGenerator[str, None]:
+    """Generator that runs speed tests at multiple sizes and yields JSON results."""
+    import json
+
+    test_sizes_mb = [10, 50, 100, 200]
+
+    for size_mb in test_sizes_mb:
+        size_bytes = size_mb * 1024 * 1024
+        result = await run_single_speed_test(size_bytes)
+
+        point = DiskSpeedTestPoint(
+            size_mb=size_mb,
+            write_speed=result.write_speed_mbps,
+            read_speed=result.read_speed_mbps,
+        )
+        yield json.dumps(point.dict())
+
+
+@disk_router.get(
+    "/speed/stream",
+    summary="Run multi-size disk speed test with streaming results.",
+)
+async def disk_speed_stream() -> StreamingResponse:
+    return StreamingResponse(
+        streamer(multi_size_speed_test_generator(), heartbeats=1.0),
+        media_type="application/x-ndjson",
+        headers={
+            "Content-Type": "application/x-ndjson",
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
 fast_api_app = FastAPI(
     title="Disk Usage API",
     description="Inspect disk usage and delete files using du.",