Add icon ID creation

rbro112 · rbro112 · commit 265a811bc364 · 2025-10-30T11:21:50.000-07:00
diff --git a/src/launchpad/artifact_processor.py b/src/launchpad/artifact_processor.py
@@ -1,5 +1,3 @@
-"""Artifact processing logic extracted from LaunchpadService."""
-
 from __future__ import annotations
 
 import contextlib
@@ -42,6 +40,7 @@
 from launchpad.size.models.apple import AppleAppInfo
 from launchpad.size.models.common import BaseAppInfo
 from launchpad.tracing import request_context
+from launchpad.utils.file_utils import IdPrefix, id_from_bytes
 from launchpad.utils.logging import get_logger
 from launchpad.utils.objectstore.service import (
     Client as ObjectstoreClient,
@@ -55,8 +54,6 @@
 
 
 class ArtifactProcessor:
-    """Handles the processing of artifacts including download, analysis, and upload."""
-
     def __init__(
         self,
         sentry_client: SentryClient,
@@ -94,7 +91,7 @@ def process_message(
         if artifact_processor is None:
             sentry_client = SentryClient(base_url=service_config.sentry_base_url)
             objectstore_client = ObjectstoreClientBuilder(
-                usecase="app-icons", options={"base_url": "http://localhost:8888/"}
+                usecase="preprod", options={"base_url": "http://localhost:8888/"}
             ).for_project(organization_id, project_id)
             artifact_processor = ArtifactProcessor(sentry_client, statsd, objectstore_client)
 
@@ -269,8 +266,11 @@ def _process_app_icon(
             logger.info(f"No app icon found for {artifact_id} (project: {project_id}, org: {organization_id})")
             return None
 
-        app_icon_id = self._objectstore_client.put(app_icon, compression="none")
-        return app_icon_id
+        image_id = id_from_bytes(app_icon, IdPrefix.ICON)
+        icon_key = f"{organization_id}/{project_id}/{image_id}"
+        logger.info(f"Uploading app icon to object store: {icon_key}")
+        self._objectstore_client.put(app_icon, compression="none", id=icon_key)
+        return image_id
 
     def _do_distribution(
         self,
diff --git a/src/launchpad/utils/file_utils.py b/src/launchpad/utils/file_utils.py
@@ -1,100 +1,78 @@
-"""File utilities for app size analyzer."""
-
 import hashlib
 import shutil
 import tempfile
 
+from enum import Enum
+from io import BytesIO
 from pathlib import Path
+from typing import IO
 
 from .logging import get_logger
 
 logger = get_logger(__name__)
 
+_HASH_CHUNK_SIZE = 8192
 
-def calculate_file_hash(file_path: Path, algorithm: str = "md5") -> str:
-    """Calculate hash of a file.
-
-    Args:
-        file_path: Path to the file
-        algorithm: Hash algorithm to use ("md5", "sha1", "sha256")
 
-    Returns:
-        Hexadecimal hash string
+class IdPrefix(Enum):
+    ICON = "icn"
+    SNAPSHOT = "snap"
 
-    Raises:
-        ValueError: If algorithm is not supported
-        FileNotFoundError: If file doesn't exist
-    """
-    if not file_path.exists():
-        raise FileNotFoundError(f"File not found: {file_path}")
 
+def _calculate_hash(data: IO[bytes], algorithm: str) -> str:
+    hasher = None
     if algorithm == "md5":
         hasher = hashlib.md5()
     elif algorithm == "sha1":
         hasher = hashlib.sha1()
     elif algorithm == "sha256":
         hasher = hashlib.sha256()
-    else:
+
+    if hasher is None:
         raise ValueError(f"Unsupported hash algorithm: {algorithm}")
 
+    for chunk in iter(lambda: data.read(_HASH_CHUNK_SIZE), b""):
+        hasher.update(chunk)
+
+    return hasher.hexdigest()
+
+
+def id_from_bytes(data: bytes, prefix: IdPrefix) -> str:
+    return f"{prefix.value}_{_calculate_hash(BytesIO(data), 'sha256')[:12]}"
+
+
+def calculate_file_hash(file_path: Path, algorithm: str = "md5") -> str:
+    if not file_path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+
     try:
         with open(file_path, "rb") as f:
-            # Read file in chunks to handle large files efficiently
-            for chunk in iter(lambda: f.read(8192), b""):
-                hasher.update(chunk)
-
-        return hasher.hexdigest()
+            return _calculate_hash(f, algorithm)
     except Exception as e:
         raise RuntimeError(f"Failed to calculate hash for {file_path}: {e}")
 
 
 def get_file_size(file_path: Path) -> int:
-    """Get file size in bytes.
-
-    Args:
-        file_path: Path to the file
-
-    Returns:
-        File size in bytes
-
-    Raises:
-        FileNotFoundError: If file doesn't exist
-    """
     if not file_path.exists():
         raise FileNotFoundError(f"File not found: {file_path}")
 
     return file_path.stat().st_size
 
 
 def to_nearest_block_size(file_size: int, block_size: int) -> int:
-    """Round file size up to the nearest filesystem block size."""
-
     if file_size == 0:
         return 0
 
     return ((file_size - 1) // block_size + 1) * block_size
 
 
 def create_temp_directory(prefix: str = "app-analyzer-") -> Path:
-    """Create a temporary directory.
-
-    Args:
-        prefix: Prefix for the temporary directory name
-
-    Returns:
-        Path to the created temporary directory
-    """
     temp_dir = Path(tempfile.mkdtemp(prefix=prefix))
     logger.debug(f"Created temporary directory: {temp_dir}")
     return temp_dir
 
 
 def cleanup_directory(directory: Path) -> None:
-    """Remove a directory and all its contents.
-
-    Args:
-        directory: Directory to remove
-    """
     if directory.exists() and directory.is_dir():
         shutil.rmtree(directory)
         logger.debug(f"Cleaned up directory: {directory}")
diff --git a/tests/unit/test_artifact_processor.py b/tests/unit/test_artifact_processor.py
@@ -1,10 +1,10 @@
-"""Tests for ArtifactProcessor including error handling, retry logic, and message processing."""
-
 from unittest.mock import Mock, patch
 
 import pytest
 
-from sentry_kafka_schemas.schema_types.preprod_artifact_events_v1 import PreprodArtifactEvents
+from sentry_kafka_schemas.schema_types.preprod_artifact_events_v1 import (
+    PreprodArtifactEvents,
+)
 
 from launchpad.artifact_processor import ArtifactProcessor
 from launchpad.constants import (
@@ -16,17 +16,17 @@
 )
 from launchpad.sentry_client import SentryClient, SentryClientError
 from launchpad.service import ServiceConfig
+from launchpad.utils.objectstore.service import Client as ObjectstoreClient
 from launchpad.utils.statsd import FakeStatsd
 
 
 class TestArtifactProcessorErrorHandling:
-    """Test error handling and retry logic in ArtifactProcessor."""
-
     def setup_method(self):
         """Set up test fixtures."""
         mock_sentry_client = Mock(spec=SentryClient)
         mock_statsd = Mock()
-        self.processor = ArtifactProcessor(mock_sentry_client, mock_statsd)
+        mock_objectstore_client = Mock(spec=ObjectstoreClient)
+        self.processor = ArtifactProcessor(mock_sentry_client, mock_statsd, mock_objectstore_client)
 
     def test_retry_operation_success_on_first_attempt(self):
         """Test that _retry_operation succeeds on first attempt."""
@@ -215,7 +215,11 @@ class TestArtifactProcessorMessageHandling:
     def test_process_message_ios(self, mock_process, mock_sentry_client):
         """Test processing iOS artifact messages."""
         fake_statsd = FakeStatsd()
-        service_config = ServiceConfig(sentry_base_url="http://test.sentry.io", projects_to_skip=[])
+        service_config = ServiceConfig(
+            sentry_base_url="http://test.sentry.io",
+            projects_to_skip=[],
+            objectstore_base_url="http://test.objectstore.io",
+        )
 
         # Create a payload for iOS artifact
         payload: PreprodArtifactEvents = {
@@ -230,20 +234,33 @@ def test_process_message_ios(self, mock_process, mock_sentry_client):
 
         # Verify process_artifact was called with correct args
         mock_process.assert_called_once_with(
-            "test-org-123", "test-project-ios", "ios-test-123", [PreprodFeature.SIZE_ANALYSIS]
+            "test-org-123",
+            "test-project-ios",
+            "ios-test-123",
+            [PreprodFeature.SIZE_ANALYSIS],
         )
 
         # Verify metrics were recorded
         calls = fake_statsd.calls
-        assert ("increment", {"metric": "artifact.processing.started", "value": 1, "tags": None}) in calls
-        assert ("increment", {"metric": "artifact.processing.completed", "value": 1, "tags": None}) in calls
+        assert (
+            "increment",
+            {"metric": "artifact.processing.started", "value": 1, "tags": None},
+        ) in calls
+        assert (
+            "increment",
+            {"metric": "artifact.processing.completed", "value": 1, "tags": None},
+        ) in calls
 
     @patch("launchpad.artifact_processor.SentryClient")
     @patch.object(ArtifactProcessor, "process_artifact")
     def test_process_message_android(self, mock_process, mock_sentry_client):
         """Test processing Android artifact messages."""
         fake_statsd = FakeStatsd()
-        service_config = ServiceConfig(sentry_base_url="http://test.sentry.io", projects_to_skip=[])
+        service_config = ServiceConfig(
+            sentry_base_url="http://test.sentry.io",
+            projects_to_skip=[],
+            objectstore_base_url="http://test.objectstore.io",
+        )
 
         # Create a payload for Android artifact
         payload: PreprodArtifactEvents = {
@@ -266,15 +283,25 @@ def test_process_message_android(self, mock_process, mock_sentry_client):
 
         # Verify metrics were recorded
         calls = fake_statsd.calls
-        assert ("increment", {"metric": "artifact.processing.started", "value": 1, "tags": None}) in calls
-        assert ("increment", {"metric": "artifact.processing.completed", "value": 1, "tags": None}) in calls
+        assert (
+            "increment",
+            {"metric": "artifact.processing.started", "value": 1, "tags": None},
+        ) in calls
+        assert (
+            "increment",
+            {"metric": "artifact.processing.completed", "value": 1, "tags": None},
+        ) in calls
 
     @patch("launchpad.artifact_processor.SentryClient")
     @patch.object(ArtifactProcessor, "process_artifact")
     def test_process_message_error(self, mock_process, mock_sentry_client):
         """Test error handling in message processing."""
         fake_statsd = FakeStatsd()
-        service_config = ServiceConfig(sentry_base_url="http://test.sentry.io", projects_to_skip=[])
+        service_config = ServiceConfig(
+            sentry_base_url="http://test.sentry.io",
+            projects_to_skip=[],
+            objectstore_base_url="http://test.objectstore.io",
+        )
 
         # Make process_artifact raise an exception
         mock_process.side_effect = RuntimeError("Download failed: HTTP 404")
@@ -292,7 +319,10 @@ def test_process_message_error(self, mock_process, mock_sentry_client):
 
         # Verify process_artifact was called
         mock_process.assert_called_once_with(
-            "test-org", "test-project", "test-123", [PreprodFeature.SIZE_ANALYSIS, PreprodFeature.BUILD_DISTRIBUTION]
+            "test-org",
+            "test-project",
+            "test-123",
+            [PreprodFeature.SIZE_ANALYSIS, PreprodFeature.BUILD_DISTRIBUTION],
         )
 
         # Verify the metrics were called correctly
@@ -308,7 +338,9 @@ def test_process_message_project_skipped(self, mock_process, mock_sentry_client)
         """Test that projects in the skip list are not processed."""
         fake_statsd = FakeStatsd()
         service_config = ServiceConfig(
-            sentry_base_url="http://test.sentry.io", projects_to_skip=["skip-project-1", "skip-project-2"]
+            sentry_base_url="http://test.sentry.io",
+            projects_to_skip=["skip-project-1", "skip-project-2"],
+            objectstore_base_url="http://test.objectstore.io",
         )
 
         # Create a payload for a project that should be skipped
@@ -334,7 +366,11 @@ def test_process_message_project_skipped(self, mock_process, mock_sentry_client)
     def test_process_message_project_not_skipped(self, mock_process, mock_sentry_client):
         """Test that projects not in the skip list are processed normally."""
         fake_statsd = FakeStatsd()
-        service_config = ServiceConfig(sentry_base_url="http://test.sentry.io", projects_to_skip=["other-project"])
+        service_config = ServiceConfig(
+            sentry_base_url="http://test.sentry.io",
+            projects_to_skip=["other-project"],
+            objectstore_base_url="http://test.objectstore.io",
+        )
 
         # Create a payload for a project that should NOT be skipped
         payload: PreprodArtifactEvents = {
@@ -357,5 +393,11 @@ def test_process_message_project_not_skipped(self, mock_process, mock_sentry_cli
 
         # Verify normal metrics were recorded
         calls = fake_statsd.calls
-        assert ("increment", {"metric": "artifact.processing.started", "value": 1, "tags": None}) in calls
-        assert ("increment", {"metric": "artifact.processing.completed", "value": 1, "tags": None}) in calls
+        assert (
+            "increment",
+            {"metric": "artifact.processing.started", "value": 1, "tags": None},
+        ) in calls
+        assert (
+            "increment",
+            {"metric": "artifact.processing.completed", "value": 1, "tags": None},
+        ) in calls
