diff --git a/src/pymax/api/uploads/service.py b/src/pymax/api/uploads/service.py index 24cdb54..7233c70 100644 --- a/src/pymax/api/uploads/service.py +++ b/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") diff --git a/tests/api/test_upload_service.py b/tests/api/test_upload_service.py index 1c1afdf..7bffafa 100644 --- a/tests/api/test_upload_service.py +++ b/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=` -- 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