fix: sort r2_uploads so it matches ids sort

Haidra-Org · Sep 26, 2024 · 58ab404 · 58ab404
1 parent 8efc611
commit 58ab404
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 3 deletions.
diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py
@@ -260,6 +260,14 @@ def ids_present(self) -> bool:
         """Whether or not the IDs are present."""
         return self._ids_present
 
+    def _sort_ids(self) -> None:
+        """Sort the IDs in place and sort so r2_uploads is changed so the same index changes occur."""
+        if len(self.ids) > 1:
+            logger.debug("Sorting IDs")
+            self.ids.sort()
+            if self.r2_uploads is not None:
+                self.r2_uploads.sort()
+
     @model_validator(mode="after")
     def validate_ids_present(self) -> ImageGenerateJobPopResponse:
         """Ensure that either id_ or ids is present."""
@@ -273,9 +281,7 @@ def validate_ids_present(self) -> ImageGenerateJobPopResponse:
         if self.id_ is None and len(self.ids) == 0:
             raise ValueError("Neither id_ nor ids were present in the response.")
 
-        if len(self.ids) > 1:
-            logger.debug("Sorting IDs")
-            self.ids.sort()
+        self._sort_ids()
 
         self._ids_present = True
 

diff --git a/tests/ai_horde_api/test_ai_horde_api_models.py b/tests/ai_horde_api/test_ai_horde_api_models.py
@@ -460,6 +460,36 @@ def test_ImageGenerateJobPopResponse() -> None:
 
     assert all(post_processor in KNOWN_UPSCALERS for post_processor in test_response.payload.post_processing)
 
+    test_response = ImageGenerateJobPopResponse(
+        ids=[
+            JobID(root=UUID("00000000-0000-0000-0000-000000000001")),
+            JobID(root=UUID("00000000-0000-0000-0000-000000000002")),
+            JobID(root=UUID("00000000-0000-0000-0000-000000000000")),
+        ],
+        payload=ImageGenerateJobPopPayload(
+            prompt="A cat in a hat",
+        ),
+        model="Deliberate",
+        r2_uploads=[
+            "https://abbaabbaabbaabbaabbaabbaabbaabba.r2.cloudflarestorage.com/horde-transient/00000000-0000-0000-0000-000000000001.webp?AWSAccessKeyId=deadbeefdeadbeefdeadbeefdeadbeef&Signature=zxcbvfakesignature%3D&Expires=1727390285",
+            "https://abbaabbaabbaabbaabbaabbaabbaabba.r2.cloudflarestorage.com/horde-transient/00000000-0000-0000-0000-000000000000.webp?AWSAccessKeyId=deadbeefdeadbeefdeadbeefdeadbeef&Signature=345567dfakes2ignature%3D&Expires=1727390285",
+            "https://abbaabbaabbaabbaabbaabbaabbaabba.r2.cloudflarestorage.com/horde-transient/00000000-0000-0000-0000-000000000002.webp?AWSAccessKeyId=deadbeefdeadbeefdeadbeefdeadbeef&Signature=asdfg32fakesignature%3D&Expires=1727390285",
+        ],
+        skipped=ImageGenerateJobPopSkippedStatus(),
+    )
+
+    assert test_response.ids_present
+    assert test_response.ids == [
+        JobID(root=UUID("00000000-0000-0000-0000-000000000000")),
+        JobID(root=UUID("00000000-0000-0000-0000-000000000001")),
+        JobID(root=UUID("00000000-0000-0000-0000-000000000002")),
+    ]
+    assert test_response.r2_uploads == [
+        "https://abbaabbaabbaabbaabbaabbaabbaabba.r2.cloudflarestorage.com/horde-transient/00000000-0000-0000-0000-000000000000.webp?AWSAccessKeyId=deadbeefdeadbeefdeadbeefdeadbeef&Signature=345567dfakes2ignature%3D&Expires=1727390285",
+        "https://abbaabbaabbaabbaabbaabbaabbaabba.r2.cloudflarestorage.com/horde-transient/00000000-0000-0000-0000-000000000001.webp?AWSAccessKeyId=deadbeefdeadbeefdeadbeefdeadbeef&Signature=zxcbvfakesignature%3D&Expires=1727390285",
+        "https://abbaabbaabbaabbaabbaabbaabbaabba.r2.cloudflarestorage.com/horde-transient/00000000-0000-0000-0000-000000000002.webp?AWSAccessKeyId=deadbeefdeadbeefdeadbeefdeadbeef&Signature=asdfg32fakesignature%3D&Expires=1727390285",
+    ]
+
 
 def test_ImageGenerateJobPopResponse_hashability() -> None:
     test_response_ids = ImageGenerateJobPopResponse(