Take the photo token by position, not from the upload URL's photoIds

src/pymax/api/uploads/service.py

@@ -4,7 +4,7 @@
 import base64
 from http import HTTPStatus
 from typing import TYPE_CHECKING
-from urllib.parse import parse_qs, quote, urlparse
+from urllib.parse import quote
 
 import aiohttp
 from pydantic import ValidationError
@@ -40,7 +40,20 @@
 
 
 class UploadService:
+    """Performs the upload half of sending an attachment.
+
+    Each ``upload_*`` method turns a local file into whatever token the matching
+    ``MSG_SEND`` opcode expects. Videos, files and voices are uploaded in chunks
+    and only complete once the server emits the corresponding processing event,
+    which is what the per-kind waiter futures below are for.
+    """
+
     def __init__(self, app: App) -> None:
+        """Subscribes to upload-completion events and opens the waiter registries.
+
+        Args:
+            app: Runtime the uploads are performed against.
+        """
         self.app = app
         self.video_upload_waiters: dict[int, asyncio.Future[VideoUploadSignal]] = {}
         self.file_upload_waiters: dict[int, asyncio.Future[FileUploadSignal]] = {}
@@ -50,6 +63,25 @@ def __init__(self, app: App) -> None:
         self.app.dispatcher.on_internal(EventType.VOICE_READY)(self.on_voice_attach)
 
     async def upload_photo(self, photo: Photo, profile: bool = False) -> AttachPhotoPayload:
+        """Uploads a single photo and returns the token to attach it with.
+
+        Requests a one-shot upload URL from MAX, POSTs the image to it, and reads
+        the resulting token. That URL carries no photo id, so the token is read
+        from the single entry of the upload result rather than looked up by id --
+        see `_extract_photo_token` for why that is exact rather than a guess.
+
+        Args:
+            photo: Image to upload.
+            profile: Whether the upload sets the account's own avatar.
+
+        Returns:
+            Attachment payload carrying the uploaded photo's token.
+
+        Raises:
+            UploadError: No upload URL was returned, the POST failed, the
+                response could not be parsed, or it did not describe exactly one
+                photo.
+        """
         logger.info("Uploading photo")
         logger.debug("Preparing photo upload payload")
 
@@ -80,19 +112,6 @@ async def upload_photo(self, photo: Photo, profile: bool = False) -> AttachPhoto
 
         logger.debug("Photo upload URL received")
 
-        try:
-            parsed_url = urlparse(url)
-            photo_id = str(parse_qs(parsed_url.query)["photoIds"][0])
-        except (KeyError, IndexError) as e:
-            logger.exception("Photo upload URL does not contain photoIds")
-            logger.debug("Invalid photo upload URL=%s", url)
-            raise UploadError("Photo upload URL does not contain photoIds") from e
-        except Exception as e:
-            logger.exception("Failed to parse photo id from upload URL")
-            logger.debug("Invalid photo upload URL=%s", url)
-            raise UploadError("Failed to parse photo id from upload URL") from e
-
-        logger.debug("Photo upload id parsed photo_id=%s", photo_id)
 
         try:
             photo_data = photo.validate_photo()
@@ -165,24 +184,38 @@ async def upload_photo(self, photo: Photo, profile: bool = False) -> AttachPhoto
             logger.debug("Invalid photo upload response=%r", result)
             raise UploadError("Invalid photo upload response model") from e
 
-        try:
-            token = model.photos[photo_id].token
-        except KeyError as e:
-            logger.exception(
-                "Photo upload response does not contain token for photo_id=%s",
-                photo_id,
+        token = self._extract_photo_token(model)
+        logger.debug("Photo upload complete")
+        return AttachPhotoPayload(photo_token=token)
+
+    @staticmethod
+    def _extract_photo_token(model: PhotoUploadResponse) -> str:
+        """Token of the single photo this request uploaded.
+
+        The upload URL no longer carries a ``photoIds`` parameter, so the token
+        is read by position instead of by photo id. That is exact rather than a
+        guess: this endpoint is always requested with ``count=1`` and MAX has
+        never batched a photo upload -- a multi-photo message is several
+        independent ``PHOTO_UPLOAD`` requests whose tokens are then sent together
+        in one ``MSG_SEND`` (see ``_upload_attachments``). The official client
+        behaves the same way.
+
+        A response that doesn't hold exactly one photo means that assumption no
+        longer holds -- there would be no way to tell which entry is ours, and
+        picking one would attach the wrong photo to a message. Fail loudly
+        instead.
+        """
+        photos = model.photos
+        if len(photos) != 1:
+            logger.error(
+                "Photo upload response holds %d photo(s), expected exactly 1: keys=%s",
+                len(photos),
+                sorted(photos),
             )
-            logger.debug("Photo upload model=%r", model)
             raise UploadError(
-                f"Photo upload response does not contain token for photo_id={photo_id}"
-            ) from e
-        except Exception as e:
-            logger.exception("Failed to extract photo token")
-            logger.debug("Photo upload model=%r", model)
-            raise UploadError("Failed to extract photo token") from e
-
-        logger.debug("Photo upload complete photo_id=%s", photo_id)
-        return AttachPhotoPayload(photo_token=token)
+                f"Photo upload response holds {len(photos)} photo(s), expected 1"
+            )
+        return next(iter(photos.values())).token
 
     async def upload_voice(self, voice: Voice) -> VoiceAttachPayload:
         logger.info("Uploading voice")

tests/api/test_upload_service.py

@@ -4,6 +4,7 @@
 
 from pymax.api.uploads.payloads import UploadPayload
 from pymax.api.uploads.service import UploadService
+from pymax.exceptions import UploadError
 from pymax.files import File, Photo, Video, Voice
 from pymax.protocol import Opcode
 from pymax.types import AttachmentType
@@ -61,7 +62,10 @@ def test_upload_payload_uses_regular_video_defaults() -> None:
 async def test_upload_photo_requests_url_posts_file_and_returns_attach_payload(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    app = FakeApp([frame({"url": "https://upload.test/path?photoIds=photo-1"})])
+    """The URL MAX returns today carries no photo id at all -- it is a one-shot
+    `uploadImage?r=<token>` -- and the result is keyed by position ("0"), so the
+    token must not be looked up by a `photoIds` query parameter."""
+    app = FakeApp([frame({"url": "https://iu.oneme.ru/uploadImage?r=TOKEN1"})])
     service = UploadService(app)
     monkeypatch.setattr(
         "pymax.api.uploads.service.aiohttp.ClientSession",
@@ -70,14 +74,115 @@ async def test_upload_photo_requests_url_posts_file_and_returns_attach_payload(
     FakeHttpSession.posts = []
     FakeHttpSession.response = FakeHttpResponse(
         200,
-        {"photos": {"photo-1": {"token": "uploaded"}}},
+        {"photos": {"0": {"token": "uploaded"}}},
     )
 
     result = await service.upload_photo(Photo(raw=b"image-bytes", name="image.jpg"))
 
     assert result.photo_token == "uploaded"
     assert app.calls[0].opcode == Opcode.PHOTO_UPLOAD
-    assert FakeHttpSession.posts[0]["url"] == "https://upload.test/path?photoIds=photo-1"
+    assert FakeHttpSession.posts[0]["url"] == "https://iu.oneme.ru/uploadImage?r=TOKEN1"
+
+
+@pytest.mark.asyncio
+async def test_upload_photo_reads_the_token_from_a_result_keyed_by_photo_id(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Keying is not guaranteed to stay "0": anything keyed works as long as the
+    response holds exactly one photo."""
+    app = FakeApp([frame({"url": "https://iu.oneme.ru/uploadImage?r=TOKEN1"})])
+    service = UploadService(app)
+    monkeypatch.setattr(
+        "pymax.api.uploads.service.aiohttp.ClientSession",
+        FakeHttpSession,
+    )
+    FakeHttpSession.posts = []
+    FakeHttpSession.response = FakeHttpResponse(
+        200,
+        {"photos": {"photo-1": {"token": "uploaded"}}},
+    )
+
+    result = await service.upload_photo(Photo(raw=b"image-bytes", name="image.jpg"))
+
+    assert result.photo_token == "uploaded"
+
+
+@pytest.mark.asyncio
+async def test_upload_photo_refuses_an_ambiguous_multi_photo_result(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Without a photo id in the URL there is no way to tell which entry is
+    ours, so a batched result must fail rather than attach the wrong photo."""
+    app = FakeApp([frame({"url": "https://iu.oneme.ru/uploadImage?r=TOKEN1"})])
+    service = UploadService(app)
+    monkeypatch.setattr(
+        "pymax.api.uploads.service.aiohttp.ClientSession",
+        FakeHttpSession,
+    )
+    FakeHttpSession.posts = []
+    FakeHttpSession.response = FakeHttpResponse(
+        200,
+        {"photos": {"0": {"token": "a"}, "1": {"token": "b"}}},
+    )
+
+    with pytest.raises(UploadError, match="expected 1"):
+        await service.upload_photo(Photo(raw=b"image-bytes", name="image.jpg"))
+
+
+@pytest.mark.asyncio
+async def test_upload_photo_refuses_an_empty_result(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """An empty result is as unusable as an ambiguous one -- there is no photo to
+    attach -- so it must fail rather than be reported as a success."""
+    app = FakeApp([frame({"url": "https://iu.oneme.ru/uploadImage?r=TOKEN1"})])
+    service = UploadService(app)
+    monkeypatch.setattr(
+        "pymax.api.uploads.service.aiohttp.ClientSession",
+        FakeHttpSession,
+    )
+    FakeHttpSession.posts = []
+    FakeHttpSession.response = FakeHttpResponse(200, {"photos": {}})
+
+    with pytest.raises(UploadError, match="expected 1"):
+        await service.upload_photo(Photo(raw=b"image-bytes", name="image.jpg"))
+
+
+@pytest.mark.asyncio
+async def test_each_photo_of_an_album_is_uploaded_separately(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A multi-photo message is N independent `count: 1` uploads -- N distinct
+    URLs, N POSTs, N results -- and the tokens are then sent together in one
+    message. The url's photo id was never what tied a result to its photo."""
+    app = FakeApp(
+        [
+            frame({"url": "https://iu.oneme.ru/uploadImage?r=TOKEN1"}),
+            frame({"url": "https://iu.oneme.ru/uploadImage?r=TOKEN2"}),
+        ]
+    )
+    service = UploadService(app)
+    monkeypatch.setattr(
+        "pymax.api.uploads.service.aiohttp.ClientSession",
+        FakeHttpSession,
+    )
+    FakeHttpSession.posts = []
+    FakeHttpSession.response = FakeHttpResponse(200, {"photos": {"0": {"token": "t"}}})
+
+    first = await service.upload_photo(Photo(raw=b"one", name="one.jpg"))
+    FakeHttpSession.response = FakeHttpResponse(200, {"photos": {"0": {"token": "t2"}}})
+    second = await service.upload_photo(Photo(raw=b"two", name="two.jpg"))
+
+    assert (first.photo_token, second.photo_token) == ("t", "t2")
+    assert [call.opcode for call in app.calls] == [Opcode.PHOTO_UPLOAD] * 2
+    assert [post["url"] for post in FakeHttpSession.posts] == [
+        "https://iu.oneme.ru/uploadImage?r=TOKEN1",
+        "https://iu.oneme.ru/uploadImage?r=TOKEN2",
+    ]
+    # One upload per photo, always: that is what makes a single-entry result
+    # unambiguous.
+    for call in app.calls:
+        assert call.payload["count"] == 1
 
 
 @pytest.mark.asyncio