From 462768b0054c4192b0d4b121946ec48c978e4fbf Mon Sep 17 00:00:00 2001
From: Ben Mares <services-git-throwaway1@tensorial.com>
Date: Wed, 11 Sep 2024 00:42:19 +0200
Subject: [PATCH] Split tests with pytest-split

---
 .github/workflows/test.yml | 84 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 81 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 2b1aca07..d3e37c0c 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -20,12 +20,59 @@ env:
 
 jobs:
   test:
-    runs-on: ${{ matrix.os }}
+    name: pytest py${{ matrix.python-version }} on ${{ matrix.os }} (${{ matrix.pytest-split-group-index }}/${{ matrix.pytest-split-group-size }})
+    runs-on: ${{ matrix.os }}-latest
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-latest, macos-latest, windows-latest ]
+        os: [ ubuntu, macos, windows ]
         python-version: [ "3.8", "3.12" ]
+
+        # pytest-split handles dividing the tests into n groups indexed 1...n.
+        # The tests are automatically split so that the expected duration of each
+        # group is roughly the same.
+        # See the "exclude" section below for pruning to the group sizes,
+        # and the "include" section for defining the group sizes.
+        pytest-split-group-index: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+
+        exclude:
+          - os: windows
+            python-version: "3.8"
+
+          # Prune the indices so that we only run the splits up to the group size
+          # defined below in "include". (This is ugly but effective.)
+          - os: ubuntu
+            pytest-split-group-index: 5
+          - os: ubuntu
+            pytest-split-group-index: 6
+          - os: ubuntu
+            pytest-split-group-index: 7
+          - os: ubuntu
+            pytest-split-group-index: 8
+          - os: ubuntu
+            pytest-split-group-index: 9
+          - os: ubuntu
+            pytest-split-group-index: 10
+          - os: macos
+            pytest-split-group-index: 5
+          - os: macos
+            pytest-split-group-index: 6
+          - os: macos
+            pytest-split-group-index: 7
+          - os: macos
+            pytest-split-group-index: 8
+          - os: macos
+            pytest-split-group-index: 9
+          - os: macos
+            pytest-split-group-index: 10
+        include:
+          - os: ubuntu
+            pytest-split-group-size: 4
+          - os: macos
+            pytest-split-group-size: 4
+          - os: windows
+            pytest-split-group-size: 10
+
     defaults:
       run:
         shell: bash -eo pipefail -l {0}
@@ -64,7 +111,10 @@ jobs:
           pytest \
             --cov=conda_lock --cov-branch --cov-report=xml --cov-report=term \
             --store-durations \
+            --clean-durations \
             --durations-path "${{ github.workspace }}/tests/durations/${{ matrix.os }}-py${{ matrix.python-version }}.json" \
+            --splits="${{ matrix.pytest-split-group-size }}" \
+            --group="${{ matrix.pytest-split-group-index }}" \
             tests
           cp coverage.xml "${{ github.workspace }}"
 
@@ -78,5 +128,33 @@ jobs:
       - name: Store test durations
         uses: actions/upload-artifact@v4
         with:
-          name: test-durations-${{ matrix.os }}-py${{ matrix.python-version }}
+          name: test-durations-${{ matrix.os }}-py${{ matrix.python-version }}-${{ matrix.pytest-split-group-index }}
           path: tests/durations/${{ matrix.os }}-py${{ matrix.python-version }}.json
+
+  aggregate-durations:
+    name: Aggregate test durations
+    runs-on: ubuntu-latest
+    needs: test
+    steps:
+      - name: Download test durations
+        uses: actions/download-artifact@v4
+      # All the artifacts are downloaded into various subdirectories.
+      # For each filename that occurs, we need to find all the files in the
+      # subdirectories with the same name, group those, and merge them.
+      - name: Construct the list of filenames to aggregate and write them to temp/filenames.txt
+        id: construct-filenames
+        run: |
+          mkdir temp
+          find . -type f -name '*.json' | xargs -n1 basename | sort | uniq > temp/filenames.txt
+          cat temp/filenames.txt
+      - name: Aggregate test durations
+        run: |
+          mkdir aggregated
+          while read -r filename; do
+            jq -s 'add' $(find . -type f -name "$filename") > "aggregated/$filename"
+          done < temp/filenames.txt
+      - name: Upload aggregated test durations
+        uses: actions/upload-artifact@v4
+        with:
+          name: aggregated-test-durations
+          path: aggregated