test_s3_and_media.py

import gc
import os
import tempfile
import unittest
from unittest.mock import MagicMock, patch

from app.services import s3_storage
from app.services.storage import Storage
from app.utils.media_url import resolve_stored_media_url


class TestS3Helpers(unittest.TestCase):
    def test_ext_to_content_type(self):
        self.assertEqual(s3_storage.ext_to_content_type("png"), "image/png")
        self.assertEqual(s3_storage.ext_to_content_type("jpeg"), "image/jpeg")
        self.assertEqual(s3_storage.ext_to_content_type("unknown"), "application/octet-stream")


class TestUploadBytes(unittest.TestCase):
    def setUp(self):
        self._bucket = os.environ.get("AWS_S3_BUCKET")
        os.environ["AWS_S3_BUCKET"] = "test-bucket-unit"

    def tearDown(self):
        if self._bucket is None:
            os.environ.pop("AWS_S3_BUCKET", None)
        else:
            os.environ["AWS_S3_BUCKET"] = self._bucket

    @patch.object(s3_storage, "_client")
    def test_upload_key_upload_prefix(self, mock_client_factory):
        mock_c = MagicMock()
        mock_client_factory.return_value = mock_c
        r = s3_storage.upload_bytes(
            prefix="upload",
            ext="png",
            body=b"\x89PNG\r\n\x1a\n",
            content_type="image/png",
        )
        self.assertEqual(r["bucket"], "test-bucket-unit")
        self.assertRegex(r["key"], r"^upload/cat_[0-9a-f]{32}\.png$")
        self.assertEqual(r["s3_uri"], f"s3://test-bucket-unit/{r['key']}")
        mock_c.put_object.assert_called_once()
        _, kwargs = mock_c.put_object.call_args
        self.assertEqual(kwargs["Bucket"], "test-bucket-unit")
        self.assertEqual(kwargs["Key"], r["key"])
        self.assertEqual(kwargs["ContentType"], "image/png")

    @patch.object(s3_storage, "_client")
    def test_upload_key_ai_photo_prefix(self, mock_client_factory):
        mock_c = MagicMock()
        mock_client_factory.return_value = mock_c
        r = s3_storage.upload_bytes(
            prefix="AI_PHOTO",
            ext="png",
            body=b"x",
            content_type="image/png",
        )
        self.assertRegex(r["key"], r"^AI_PHOTO/cat_[0-9a-f]{32}\.png$")


class TestPresignAndMediaUrl(unittest.TestCase):
    @patch.object(s3_storage, "_client")
    def test_presigned_get_url(self, mock_client_factory):
        mock_c = MagicMock()
        mock_c.generate_presigned_url.return_value = "https://example-signed"
        mock_client_factory.return_value = mock_c
        url = s3_storage.presigned_get_url("s3://myb/path/to/cat_abc.png", expires=60)
        self.assertEqual(url, "https://example-signed")
        mock_c.generate_presigned_url.assert_called_once()

    def test_resolve_stored_media_url_passes_through(self):
        cfg = {"AWS_S3_PRESIGN_TTL": 3600}
        self.assertIsNone(resolve_stored_media_url(cfg, None))
        self.assertEqual(
            resolve_stored_media_url(cfg, "/uploads/cat_x.png"),
            "/uploads/cat_x.png",
        )

    @patch("app.utils.media_url.presigned_get_url")
    def test_resolve_s3_delegates_to_presign(self, mock_presign):
        mock_presign.return_value = "https://signed-test"
        cfg = {"AWS_S3_PRESIGN_TTL": 7200}
        out = resolve_stored_media_url(cfg, "s3://b/upload/cat_1.png")
        self.assertEqual(out, "https://signed-test")
        mock_presign.assert_called_once_with(
            "s3://b/upload/cat_1.png",
            expires=7200,
        )


class TestStorageMediaMigration(unittest.TestCase):
    def test_ai_column_and_offspring_table(self):
        fd, path = tempfile.mkstemp(suffix=".db")
        os.close(fd)
        try:
            store = Storage(path)
            with store._conn() as conn:
                cols = {r[1] for r in conn.execute("PRAGMA table_info(cats)").fetchall()}
                self.assertIn("ai_match_photo_s3", cols)
                tables = {
                    r[0]
                    for r in conn.execute(
                        "SELECT name FROM sqlite_master WHERE type='table'"
                    ).fetchall()
                }
                self.assertIn("offspring_ai_photos", tables)
            store.update_cat_ai_match_photo(1, "s3://b/AI_PHOTO/cat_x.png")
            store.upsert_offspring_ai_photo(1, 99, "s3://b/AI_PHOTO/cat_y.png")
            with store._conn() as conn:
                row = conn.execute(
                    "SELECT s3_path FROM offspring_ai_photos WHERE user_id=1 AND match_cat_id=99"
                ).fetchone()
                self.assertEqual(row[0], "s3://b/AI_PHOTO/cat_y.png")
        finally:
            del store
            gc.collect()
            try:
                os.unlink(path)
            except PermissionError:
                pass


if __name__ == "__main__":
    unittest.main()