Add query limit test and pipeline collateral

concourse/pipeline.yml

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

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-python-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-python-repository/tests/docker-compose-limits-tests.yml run query-limits-tests
+      # stop and remove containers
+      docker-compose -f fauna-python-repository/tests/docker-compose-limits-tests.yml down
+      docker-compose -f testtools-repo/fauna-driver-query-limits-tests/docker-compose.yml down

tests/docker-compose-limits-tests.yml

@@ -0,0 +1,22 @@
+version: "3.5"
+
+networks:
+  limit-net:
+    external: true
+    name: limit-net
+
+services:
+  query-limits-tests:
+    environment:
+      FAUNA_ENDPOINT: ${FAUNA_ENDPOINT:-http://fauna-limits:8443}
+      QUERY_LIMITS_DB: ${QUERY_LIMITS_DB}
+      QUERY_LIMITS_COLL: ${QUERY_LIMITS_COLL}
+    image: python:3.12.0-alpine3.17
+    container_name: query-limits-tests
+    networks:
+      - limit-net
+    volumes:
+      - "../:/tmp/app"
+    working_dir: "/tmp/app"
+    command:
+      - tests/run-query-limits-tests.sh

tests/integration/test_client_with_query_limits.py

@@ -0,0 +1,41 @@
+from multiprocessing.pool import ThreadPool
+import os
+
+import pytest
+
+from fauna import fql
+from fauna.client import Client
+from fauna.encoding import QuerySuccess
+
+
+def query_collection(client: Client) -> QuerySuccess:
+  coll_name = os.environ.get("QUERY_LIMITS_COLL")
+  return client.query(fql("${coll}.all().paginate(50)", coll=fql(coll_name)))
+
+
+@pytest.mark.skipif("QUERY_LIMITS_DB" not in os.environ or "QUERY_LIMITS_COLL" not in os.environ,
+                    reason="QUERY_LIMITS_DB and QUERY_LIMITS_COLL must both be set to run this test")
+def test_client_retries_throttled_query():
+  db_name = os.environ.get("QUERY_LIMITS_DB")
+  rootClient = Client()
+  res = rootClient.query(fql("""
+if (Database.byName(${db}).exists()) {
+  Key.create({ role: "admin", database: ${db} }) { secret }
+} else {
+  abort("Database not found.")
+}""", db=db_name))
+  secret = res.data["secret"]
+  clients = [ Client(secret=secret) for _ in range(5) ]
+  throttled = False
+
+  with ThreadPool() as pool:
+    results = pool.map(query_collection, clients)
+
+    for result in results:
+      if result.stats.attempts > 1:
+        throttled = True
+
+  assert throttled == True
+
+  for cl in clients:
+    cl.close()

tests/run-query-limits-tests.sh

@@ -0,0 +1,11 @@
+#!/bin/sh
+# NOTE: this file is intended to be ran from within the Docker image
+set -eou pipefail
+
+apk add --update make
+
+pip install . ".[test]"
+pip install coverage
+
+python -m coverage run -m pytest -v tests/integration/test_client_with_query_limits.py
+python -m coverage report -m