BT-4181: Add test for query-limited account (#129)

client_query_limits_test.go

@@ -0,0 +1,87 @@
+package fauna_test
+
+import (
+	"os"
+	"sync"
+	"testing"
+
+	"github.com/fauna/fauna-go"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestClientRetriesWithQueryLimits(t *testing.T) {
+	t.Run("Query limits succeed on retry", func(t *testing.T) {
+		dbName, dbNameSet := os.LookupEnv("QUERY_LIMITS_DB")
+		collName, collNameSet := os.LookupEnv("QUERY_LIMITS_COLL")
+
+		// If run in a pipeline, these will be empty strings, so check both
+		if (!dbNameSet || !collNameSet) ||
+			(dbName == "" || collName == "") {
+			t.Skip("Skipping query limits test due to missing env var(s)")
+		}
+
+		if _, found := os.LookupEnv(fauna.EnvFaunaSecret); !found {
+			t.Setenv(fauna.EnvFaunaSecret, "secret")
+		}
+
+		client, clientErr := fauna.NewDefaultClient()
+		if !assert.NoError(t, clientErr) {
+			return
+		}
+
+		type secretObj struct {
+			Secret string `fauna:"secret"`
+		}
+
+		query, _ := fauna.FQL(`
+		if (Database.byName(${dbName}).exists()) {
+      Key.create({ role: "admin", database: ${dbName} }) { secret }
+    } else {
+      abort("Database not found.")
+    }`, map[string]any{"dbName": dbName})
+
+		res, queryErr := client.Query(query)
+		if !assert.NoError(t, queryErr) {
+			t.FailNow()
+		}
+
+		var secret secretObj
+		marshalErr := res.Unmarshal(&secret)
+		if assert.NoError(t, marshalErr) {
+			clients := make([]*fauna.Client, 5)
+			results := make(chan int, len(clients))
+
+			var wg sync.WaitGroup
+			wg.Add(len(clients))
+
+			for i := range clients {
+				clients[i] = fauna.NewClient(secret.Secret, fauna.DefaultTimeouts(), fauna.URL(os.Getenv(fauna.EnvFaunaEndpoint)))
+
+				go func(collName string, client *fauna.Client, result chan int) {
+					defer wg.Done()
+					coll, _ := fauna.FQL(collName, nil)
+					q, _ := fauna.FQL(`${coll}.all().paginate(50)`, map[string]any{"coll": coll})
+					res, err := client.Query(q)
+					if err != nil {
+						result <- -1
+					} else {
+						result <- res.Stats.Attempts
+					}
+				}(collName, clients[i], results)
+			}
+
+			go func() {
+				wg.Wait()
+				close(results)
+			}()
+
+			throttled := false
+
+			for result := range results {
+				throttled = throttled || result > 1
+			}
+
+			assert.True(t, throttled)
+		}
+	})
+}

concourse/pipeline.yml

@@ -89,6 +89,13 @@ jobs:
                 FAUNA_SECRET: ((drivers-platform-tests/fauna-secret))
                 VERCEL_TOKEN: ((drivers-platform-tests/vercel-token))
 
+            - task: query-limits-tests
+              privileged: true
+              file: fauna-go-repository/concourse/tasks/query-limits-tests.yml
+              params:
+                QUERY_LIMITS_DB: limited
+                QUERY_LIMITS_COLL: limitCollection
+
 
   - name: release
     serial: true

concourse/scripts/docker-compose-fauna-limits.yml

@@ -0,0 +1,25 @@
+version: "3.8"
+
+networks:
+  limit-net:
+    external: true
+    name: limit-net
+
+services:
+  query-limits-tests:
+    image: golang:1.21-alpine
+    entrypoint: ["/bin/sh", "-c"]
+    volumes:
+      - "../..:/tmp/app"
+    working_dir: "/tmp/app"
+    environment:
+      FAUNA_ENDPOINT: ${FAUNA_ENDPOINT:-http://fauna-limits:8443}
+      QUERY_LIMITS_DB: ${QUERY_LIMITS_DB}
+      QUERY_LIMITS_COLL: ${QUERY_LIMITS_COLL}
+    networks:
+      - limit-net
+    command:
+    - |
+      apk add build-base
+      CGO_ENABLED=0 go build .
+      go test -v client_query_limits_test.go

concourse/tasks/query-limits-tests.yml

@@ -0,0 +1,32 @@
+---
+platform: linux
+image_resource:
+  type: registry-image
+  source:
+    repository: shared-concourse-dind
+    aws_access_key_id: ((prod-images-aws-access-key-id))
+    aws_secret_access_key: ((prod-images-aws-secret-key))
+    aws_region: us-east-2
+
+params:
+  FAUNA_ENDPOINT: http://fauna-limits:8443
+  QUERY_LIMITS_DB:
+  QUERY_LIMITS_COLL:
+
+inputs:
+  - name: fauna-go-repository
+  - name: testtools-repo
+
+run:
+  path: entrypoint.sh
+  args:
+    - bash
+    - -ceu
+    - |
+      # setup Fauna container
+      docker-compose -f testtools-repo/fauna-driver-query-limits-tests/docker-compose.yml run setup
+      # run tests
+      docker-compose -f fauna-go-repository/concourse/scripts/docker-compose-fauna-limits.yml run query-limits-tests
+      # stop and remove containers
+      docker-compose -f fauna-go-repository/concourse/scripts/docker-compose-fauna-limits.yml down
+      docker-compose -f testtools-repo/fauna-driver-query-limits-tests/docker-compose.yml down