From 5c0a252e63a4920804ab7b4d81f16226be088df4 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Wed, 11 Sep 2024 15:34:59 +0200
Subject: [PATCH] Fingers crossed.

---
 .github/workflows/build.yaml                         |  2 --
 integration-tests/conftest.py                        |  3 +++
 .../test_flash_llama_completion_many_prompts.json    | 12 ++++++------
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 6a543488..83648640 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -181,8 +181,6 @@ jobs:
       - name: Install
         run: |
           make install-integration-tests
-      - name: Setup tmate session
-        uses: mxschmitt/action-tmate@v3
       - name: Run tests
         run: |
           export DOCKER_VOLUME=/mnt/cache
diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
index 7f636373..9632c816 100644
--- a/integration-tests/conftest.py
+++ b/integration-tests/conftest.py
@@ -103,6 +103,9 @@ class ResponseComparator(JSONSnapshotExtension):
         def _convert_data(data):
             if isinstance(data, Dict):
                 if "choices" in data:
+                    data["choices"] = list(
+                        sorted(data["choices"], key=lambda x: x["index"])
+                    )
                     choices = data["choices"]
                     if isinstance(choices, List) and len(choices) >= 1:
                         if "delta" in choices[0]:
diff --git a/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts.json b/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts.json
index abbc29ac..25b8120d 100644
--- a/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts.json
+++ b/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts.json
@@ -1,5 +1,11 @@
 {
   "choices": [
+    {
+      "finish_reason": "length",
+      "index": 0,
+      "logprobs": null,
+      "text": " A Beginner’s Guide\nDeep learning is a subset"
+    },
     {
       "finish_reason": "length",
       "index": 1,
@@ -12,12 +18,6 @@
       "logprobs": null,
       "text": "usculas_minusculas(s):\n    \"\"\"\n"
     },
-    {
-      "finish_reason": "length",
-      "index": 0,
-      "logprobs": null,
-      "text": " A Beginner’s Guide\nDeep learning is a subset"
-    },
     {
       "finish_reason": "length",
       "index": 2,