ci: split tests into chunks

We have lots of tests and would like to add more. However, adding more
testing coverage comes at the cost of a slower CI when we run them
sequentially. This adds test outputs that are chunked however we'd like
to tune for batch sizes. Allowing us to create a parallelized CI
workflow.

Signed-off-by: Austin Horstman <khaneliman12@gmail.com>

.github/workflows/test.yml

@@ -39,12 +39,37 @@ jobs:
               - 'flake.lock'
               - 'flake.nix'
               - 'home-manager/**'
+  get-test-chunks:
+    runs-on: ubuntu-latest
+    outputs:
+      test-matrix: ${{ steps.chunks.outputs.test-matrix }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: cachix/install-nix-action@v31
+        with:
+          extra_nix_config: |
+            experimental-features = nix-command flakes
+      - name: Get test chunks
+        id: chunks
+        run: |
+          linux_chunks=$(nix eval --json ./tests#packages.x86_64-linux --apply 'pkgs: builtins.attrNames (builtins.removeAttrs pkgs ["metadata"])' | jq -r '[.[] | select(startswith("test-chunk-")) | sub("test-chunk-"; "") | tonumber] | sort')
+          echo "Found Linux chunks: $linux_chunks"
+
+          darwin_chunks=$(nix eval --json ./tests#packages.aarch64-darwin --apply 'pkgs: builtins.attrNames (builtins.removeAttrs pkgs ["metadata"])' | jq -r '[.[] | select(startswith("test-chunk-")) | sub("test-chunk-"; "") | tonumber] | sort')
+          echo "Found Darwin chunks: $darwin_chunks"
+
+          matrix=$(jq -n -c \
+            --argjson linux_chunks "$linux_chunks" \
+            --argjson darwin_chunks "$darwin_chunks" \
+            '{include: [($linux_chunks[] | {os: "ubuntu-latest", test_chunk: .}), ($darwin_chunks[] | {os: "macos-latest", test_chunk: .})]}')
+
+          echo "test-matrix=$matrix" >> $GITHUB_OUTPUT
+          echo "Generated matrix: $matrix"
   tests:
-    needs: changes
+    needs: [changes, get-test-chunks]
     strategy:
       fail-fast: false
-      matrix:
-        os: [ubuntu-latest, macos-latest]
+      matrix: ${{ fromJson(needs.get-test-chunks.outputs.test-matrix) }}
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v4
@@ -58,15 +83,17 @@ jobs:
           nix_path: nixpkgs=https://github.com/NixOS/nixpkgs/archive/${{ steps.get-nixpkgs.outputs.rev }}.tar.gz
           extra_nix_config: |
             experimental-features = nix-command flakes
+            max-jobs = auto
+            cores = 0
       - name: Build docs
-        if: github.event_name == 'schedule' || needs.changes.outputs.docs == 'true'
+        if: matrix.os == 'ubuntu-latest' && matrix.test_chunk == 1 && (github.event_name == 'schedule' || needs.changes.outputs.docs == 'true')
         run: nix build --show-trace .#docs-jsonModuleMaintainers
       - name: Format Check
-        if: github.event_name == 'schedule' || needs.changes.outputs.format == 'true'
+        if: matrix.os == 'ubuntu-latest' && matrix.test_chunk == 1 && (github.event_name == 'schedule' || needs.changes.outputs.format == 'true')
         run: nix fmt -- --ci
       - name: Test init --switch with locked inputs
         # FIXME: nix broken on darwin on unstable
-        if: matrix.os != 'macos-latest' && (github.event_name == 'schedule' || needs.changes.outputs.hm == 'true')
+        if: matrix.os == 'ubuntu-latest' && matrix.test_chunk == 1 && (github.event_name == 'schedule' || needs.changes.outputs.hm == 'true')
         run: |
           # Copy lock file to home directory for consistent testing
           mkdir -p ~/.config/home-manager
@@ -74,16 +101,18 @@ jobs:
           nix run .#home-manager -- init --switch --override-input home-manager .
       - name: Uninstall
         # FIXME: nix broken on darwin on unstable
-        if: matrix.os != 'macos-latest' && (github.event_name == 'schedule' || needs.changes.outputs.hm == 'true')
+        if: matrix.os == 'ubuntu-latest' && matrix.test_chunk == 1 && (github.event_name == 'schedule' || needs.changes.outputs.hm == 'true')
         run: yes | nix run . -- uninstall
-      - name: Run tests
+      - name: Run tests (chunk ${{ matrix.test_chunk }})
         if: github.event_name == 'schedule' || needs.changes.outputs.tests == 'true'
-        run: nix build -j auto --show-trace --option allow-import-from-derivation false --reference-lock-file flake.lock "./tests#test-all-no-big"
+        run: |
+          nix build -j auto --show-trace --option allow-import-from-derivation false --reference-lock-file flake.lock "./tests#test-chunk-${{ matrix.test_chunk }}"
         env:
           GC_INITIAL_HEAP_SIZE: 4294967296
-      - name: Run tests (with IFD)
+      - name: Run tests with IFD (chunk ${{ matrix.test_chunk }})
         if: github.event_name == 'schedule' || needs.changes.outputs.tests == 'true'
-        run: nix build -j auto --show-trace --reference-lock-file flake.lock "./tests#test-all-no-big"
+        run: |
+          nix build -j auto --show-trace --reference-lock-file flake.lock "./tests#test-chunk-${{ matrix.test_chunk }}"
         env:
           GC_INITIAL_HEAP_SIZE: 4294967296
       - name: Generate Job Summary

tests/flake.nix

@@ -70,9 +70,47 @@
               };
             in
             lib.nameValuePair "test-all-no-big-ifd" tests.build.all;
+
+          # Create chunked test packages for better CI parallelization
+          testChunks =
+            let
+              tests = import ./. {
+                inherit pkgs;
+                # Disable big tests since this is only used for CI
+                enableBig = false;
+              };
+              allTests = lib.attrNames tests.build;
+              # Remove 'all' from the test list as it's a meta-package
+              filteredTests = lib.filter (name: name != "all") allTests;
+              # NOTE: Just a starting value, we can tweak this to find a good value.
+              targetTestsPerChunk = 250;
+              numChunks = lib.max 1 (
+                builtins.ceil ((builtins.length filteredTests) / (targetTestsPerChunk * 1.0))
+              );
+              chunkSize = builtins.ceil ((builtins.length filteredTests) / (numChunks * 1.0));
+
+              makeChunk =
+                chunkNum: testList:
+                let
+                  start = (chunkNum - 1) * chunkSize;
+                  end = lib.min (start + chunkSize) (builtins.length testList);
+                  chunkTests = lib.sublist start (end - start) testList;
+                  chunkAttrs = lib.genAttrs chunkTests (name: tests.build.${name});
+                in
+                pkgs.symlinkJoin {
+                  name = "test-chunk-${toString chunkNum}";
+                  paths = lib.attrValues chunkAttrs;
+                };
+            in
+            lib.listToAttrs (
+              lib.genList (
+                i: lib.nameValuePair "test-chunk-${toString (i + 1)}" (makeChunk (i + 1) filteredTests)
+              ) numChunks
+            );
         in
         testPackages
         // integrationTestPackages
+        // testChunks
         // (lib.listToAttrs [
           testAllNoBig
           testAllNoBigIfd