soleprint/station/tools/tester/playwright/artifacts.py

"""
Artifact storage and retrieval for test results.
"""

import shutil
from pathlib import Path
from typing import Optional
from dataclasses import dataclass


@dataclass
class TestArtifact:
    """Test artifact (video, screenshot, trace, etc.)."""
    type: str  # "video", "screenshot", "trace", "log"
    filename: str
    path: str
    size: int
    mimetype: str
    url: str  # Streaming endpoint


class ArtifactStore:
    """Manage test artifacts."""

    def __init__(self, artifacts_dir: Path):
        self.artifacts_dir = artifacts_dir
        self.videos_dir = artifacts_dir / "videos"
        self.screenshots_dir = artifacts_dir / "screenshots"
        self.traces_dir = artifacts_dir / "traces"

        # Ensure directories exist
        self.videos_dir.mkdir(parents=True, exist_ok=True)
        self.screenshots_dir.mkdir(parents=True, exist_ok=True)
        self.traces_dir.mkdir(parents=True, exist_ok=True)

    def store_artifact(
        self,
        source_path: Path,
        run_id: str,
        artifact_type: str
    ) -> Optional[TestArtifact]:
        """
        Store an artifact and return its metadata.

        Args:
            source_path: Path to the source file
            run_id: Test run ID
            artifact_type: Type of artifact (video, screenshot, trace)

        Returns:
            TestArtifact metadata or None if storage fails
        """
        if not source_path.exists():
            return None

        # Determine destination directory
        if artifact_type == "video":
            dest_dir = self.videos_dir
            mimetype = "video/webm"
        elif artifact_type == "screenshot":
            dest_dir = self.screenshots_dir
            mimetype = "image/png"
        elif artifact_type == "trace":
            dest_dir = self.traces_dir
            mimetype = "application/zip"
        else:
            # Unknown type, store in root artifacts dir
            dest_dir = self.artifacts_dir
            mimetype = "application/octet-stream"

        # Create run-specific subdirectory
        run_dir = dest_dir / run_id
        run_dir.mkdir(parents=True, exist_ok=True)

        # Copy file
        dest_path = run_dir / source_path.name
        try:
            shutil.copy2(source_path, dest_path)
        except Exception:
            return None

        # Build streaming URL
        url = f"/tools/tester/api/artifact/{run_id}/{source_path.name}"

        return TestArtifact(
            type=artifact_type,
            filename=source_path.name,
            path=str(dest_path),
            size=dest_path.stat().st_size,
            mimetype=mimetype,
            url=url,
        )

    def get_artifact(self, run_id: str, filename: str) -> Optional[Path]:
        """
        Retrieve an artifact file.

        Args:
            run_id: Test run ID
            filename: Artifact filename

        Returns:
            Path to artifact file or None if not found
        """
        # Search in all artifact directories
        for artifact_dir in [self.videos_dir, self.screenshots_dir, self.traces_dir]:
            artifact_path = artifact_dir / run_id / filename
            if artifact_path.exists():
                return artifact_path

        # Check root artifacts dir
        artifact_path = self.artifacts_dir / run_id / filename
        if artifact_path.exists():
            return artifact_path

        return None

    def list_artifacts(self, run_id: str) -> list[TestArtifact]:
        """
        List all artifacts for a test run.

        Args:
            run_id: Test run ID

        Returns:
            List of TestArtifact metadata
        """
        artifacts = []

        # Search in all artifact directories
        type_mapping = {
            self.videos_dir: ("video", "video/webm"),
            self.screenshots_dir: ("screenshot", "image/png"),
            self.traces_dir: ("trace", "application/zip"),
        }

        for artifact_dir, (artifact_type, mimetype) in type_mapping.items():
            run_dir = artifact_dir / run_id
            if not run_dir.exists():
                continue

            for artifact_file in run_dir.iterdir():
                if artifact_file.is_file():
                    artifacts.append(TestArtifact(
                        type=artifact_type,
                        filename=artifact_file.name,
                        path=str(artifact_file),
                        size=artifact_file.stat().st_size,
                        mimetype=mimetype,
                        url=f"/tools/tester/api/artifact/{run_id}/{artifact_file.name}",
                    ))

        return artifacts

    def cleanup_old_artifacts(self, keep_recent: int = 10):
        """
        Clean up old artifact directories, keeping only the most recent runs.

        Args:
            keep_recent: Number of recent runs to keep
        """
        # Get all run directories sorted by modification time
        all_runs = []

        for artifact_dir in [self.videos_dir, self.screenshots_dir, self.traces_dir]:
            for run_dir in artifact_dir.iterdir():
                if run_dir.is_dir():
                    all_runs.append(run_dir)

        # Sort by modification time (newest first)
        all_runs.sort(key=lambda p: p.stat().st_mtime, reverse=True)

        # Keep only the most recent
        for old_run in all_runs[keep_recent:]:
            try:
                shutil.rmtree(old_run)
            except Exception:
                pass  # Ignore errors during cleanup