From d9adaa16c5bcde9e1306a9a8ba5e0abaa4435ec8 Mon Sep 17 00:00:00 2001
From: Alec Thomas <aat@tbd.email>
Date: Tue, 17 Dec 2024 19:24:26 +1100
Subject: [PATCH] chore(ci): multiple integration tests per shard in CI

---
 .github/workflows/ci.yml        |  4 ++--
 Justfile                        | 13 +++----------
 internal/integration/actions.go |  3 +--
 scripts/retry                   | 21 +++++++++++++++++++++
 4 files changed, 27 insertions(+), 14 deletions(-)
 create mode 100755 scripts/retry

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index c064bdae4..d75f86a04 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -257,7 +257,7 @@ jobs:
         run: |
           set -euo pipefail
           # shellcheck disable=SC2046
-          echo "matrix={\"test\":$(jq -c -n '$ARGS.positional' --args $(git grep -l '^//go:build integration' | xargs grep '^func Test' | awk '{print $2}' | cut -d'(' -f1))}" >> "$GITHUB_OUTPUT"
+          echo "matrix={\"test\":$(jq -c -n '$ARGS.positional' --args $(git grep -l '^//go:build integration' | xargs grep '^func Test' | awk '{print $2}' | cut -d'(' -f1 | paste -d '|' - - | sed 's/|*$//'))}" >> "$GITHUB_OUTPUT"
   integration-run:
     name: Integration Test
     # if: github.event_name != 'pull_request' || github.event.action == 'enqueued' || contains( github.event.pull_request.labels.*.name, 'run-all')
@@ -281,7 +281,7 @@ jobs:
         run: |
           set -euo pipefail
           # shellcheck disable=SC2046
-          go test -v -race -tags integration -run '^${{ matrix.test }}$' $(git grep -l '^//go:build integration' | xargs grep -l '^func ${{ matrix.test }}' | xargs -I {} dirname ./{})
+          echo '${{ matrix.test }}' | tr '|' ' ' | xargs -n1 just integration-tests
   infrastructure-shard:
     name: Shard Infrastructure Tests
     # if: github.event_name != 'pull_request' || github.event.action == 'enqueued' || contains( github.event.pull_request.labels.*.name, 'run-all')
diff --git a/Justfile b/Justfile
index 830cf9167..f194da596 100644
--- a/Justfile
+++ b/Justfile
@@ -205,9 +205,7 @@ format-frontend:
 
 # Install Node dependencies using pnpm
 pnpm-install:
-  #!/bin/bash
-  set -euo pipefail
-  for i in {1..3}; do mk frontend/{console,vscode}/node_modules : frontend/{console,vscode}/package.json -- "pnpm install --frozen-lockfile" && break || sleep 5; done
+  retry 3 pnpm install --frozen-lockfile
 
 # Copy plugin protos from the SQLC release
 update-sqlc-plugin-codegen-proto:
@@ -233,16 +231,11 @@ build-protos-unconditionally: go2proto lint-protos pnpm-install
 # Run integration test(s)
 integration-tests *test:
   #!/bin/bash
-  set -euo pipefail
-  testName=${1:-}
-  for i in {1..3}; do go test -fullpath -count 1 -v -tags integration -run "$testName" -p 1 $(find . -type f -name '*_test.go' -print0 | xargs -0 grep -r -l "$testName" | xargs grep -l '//go:build integration' | xargs -I {} dirname './{}') && break || true; done
+  retry 3 /bin/bash -c "go test -fullpath -count 1 -v -tags integration -run '^({{test}})$' -p 1 $(find . -type f -name '*_test.go' -print0 | xargs -0 grep -r -l {{test}} | xargs grep -l '//go:build integration' | xargs -I {} dirname './{}' | tr '\n' ' ')"
 
 # Run integration test(s)
 infrastructure-tests *test:
-  #!/bin/bash
-  set -euo pipefail
-  testName=${1:-}
-  for i in {1..3}; do go test -fullpath -count 1 -v -tags infrastructure -run "$testName" -p 1 $(find . -type f -name '*_test.go' -print0 | xargs -0 grep -r -l "$testName" | xargs grep -l '//go:build infrastructure' | xargs -I {} dirname './{}') && break || true; done
+  retry 3 /bin/bash -c "go test -fullpath -count 1 -v -tags infrastructure -run '^({{test}})$' -p 1 $(find . -type f -name '*_test.go' -print0 | xargs -0 grep -r -l {{test}} | xargs grep -l '//go:build infrastructure' | xargs -I {} dirname './{}' | tr '\n' ' ')"
 
 # Run README doc tests
 test-readme *args:
diff --git a/internal/integration/actions.go b/internal/integration/actions.go
index bc2c06756..bd1bbd4c4 100644
--- a/internal/integration/actions.go
+++ b/internal/integration/actions.go
@@ -7,7 +7,6 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
-	"fmt"
 	"io"
 	"net/http"
 	"net/url"
@@ -624,7 +623,7 @@ func JsonData(t testing.TB, body interface{}) []byte {
 func HttpCall(method string, path string, headers map[string][]string, body []byte, onResponse func(t testing.TB, resp *HTTPResponse)) Action {
 	return func(t testing.TB, ic TestContext) {
 		Infof("HTTP %s %s", method, path)
-		baseURL, err := url.Parse(fmt.Sprintf("http://localhost:8891"))
+		baseURL, err := url.Parse("http://localhost:8891")
 		assert.NoError(t, err)
 
 		u, err := baseURL.Parse(path)
diff --git a/scripts/retry b/scripts/retry
new file mode 100755
index 000000000..5a69ab511
--- /dev/null
+++ b/scripts/retry
@@ -0,0 +1,21 @@
+#!/bin/bash
+# #
+# Retry a command until it succeeds or the maximum number of retries is reached.
+# Usage: retry <max-retries> <command>
+
+set -euo pipefail
+
+max_retries=${1:-}
+shift
+
+for i in $(seq 1 "$max_retries"); do
+  echo "$@"
+  if "$@"; then
+    exit 0
+  fi
+  if [ "$i" -eq "$max_retries" ]; then
+    exit 1
+  fi
+  echo "Failed, retrying in 5 seconds..."
+  sleep 5
+done