From 4d943425a0339020f369e27bcbf06858fc44965a Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Wed, 13 Nov 2024 10:59:50 -0500
Subject: [PATCH] add task

---
 .evergreen/config.yml                         | 57 +++++++++++++++++++
 Taskfile.yml                                  |  3 +
 .../client_side_encryption_prose_test.go      |  7 ++-
 3 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/.evergreen/config.yml b/.evergreen/config.yml
index 90da768350..d2bf4b101b 100644
--- a/.evergreen/config.yml
+++ b/.evergreen/config.yml
@@ -553,6 +553,42 @@ functions:
           KMS_MOCK_SERVERS_RUNNING: "true"
         args: [*task-runner, evg-test-kmip]
 
+  start-kms-failpoint-server:
+    - command: ec2.assume_role
+      params:
+        role_arn: ${aws_test_secrets_role}
+    - command: subprocess.exec
+      params:
+        working_dir: src/go.mongodb.org/mongo-driver
+        binary: bash
+        background: true
+        include_expansions_in_env: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "DRIVERS_TOOLS"]
+        # This cannot use task because it will hang on Windows.
+        args: [etc/setup-encryption.sh]
+    - command: subprocess.exec
+      params:
+        binary: python3
+        background: true
+        args: ["-u", "${DRIVERS_TOOLS}/.evergreen/csfle/kms_failpoint_server.py", "--port", "9003"]
+
+  run-retry-kms-requests:
+    - command: subprocess.exec
+      type: test
+      params:
+        binary: "bash"
+        env: 
+          GO_BUILD_TAGS: cse
+        include_expansions_in_env: [AUTH, SSL, MONGODB_URI, TOPOLOGY, 
+          MONGO_GO_DRIVER_COMPRESSOR]
+        args: [*task-runner, setup-test]
+    - command: subprocess.exec
+      type: test
+      params:
+        binary: "bash"
+        env:
+          KMS_FAILPOINT_SERVERS_RUNNING: "true"
+        args: [*task-runner, evg-test-retry-kms-requests]
+
   run-fuzz-tests:
     - command: subprocess.exec
       type: test
@@ -1486,6 +1522,21 @@ tasks:
           AUTH: "noauth"
           SSL: "nossl"
 
+  - name: "test-retry-kms-requests"
+    tags: ["retry-kms-requests"]
+    commands:
+      - func: bootstrap-mongo-orchestration
+        vars:
+          TOPOLOGY: "server"
+          AUTH: "noauth"
+          SSL: "nossl"
+      - func: start-kms-failpoint-server
+      - func: run-retry-kms-requests
+        vars:
+          TOPOLOGY: "server"
+          AUTH: "noauth"
+          SSL: "nossl"
+
   - name: "test-serverless"
     tags: ["serverless"]
     commands:
@@ -2195,6 +2246,12 @@ buildvariants:
     tasks:
       - name: ".kms-kmip"
 
+  - matrix_name: "retry-kms-requests-test"
+    matrix_spec: { version: ["7.0"], os-ssl-40: ["rhel87-64"] }
+    display_name: "Retry KMS Requests ${os-ssl-40}"
+    tasks:
+      - name: ".retry-kms-requests"
+
   - matrix_name: "fuzz-test"
     matrix_spec: { version: ["5.0"], os-ssl-40: ["rhel87-64"] }
     display_name: "Fuzz ${version} ${os-ssl-40}"
diff --git a/Taskfile.yml b/Taskfile.yml
index f22427a640..d93a7e4e9b 100644
--- a/Taskfile.yml
+++ b/Taskfile.yml
@@ -141,6 +141,9 @@ tasks:
   evg-test-kms:
     - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_tls_tests >> test.suite
 
+  evg-test-retry-kms-requests:
+    - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_retry_tests >> test.suite
+
   evg-test-load-balancers:
     # Load balancer should be tested with all unified tests as well as tests in the following
     # components: retryable reads, retryable writes, change streams, initial DNS seedlist discovery.
diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index 4b6966c47b..604d025f8c 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -2983,7 +2983,12 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		})
 	})
 
-	mt.RunOpts("24. KMS Retry Tests", qeRunOpts22, func(mt *mtest.T) {
+	mt.RunOpts("24. kms retry tests", noClientOpts, func(mt *mtest.T) {
+		kmsTlsTestcase := os.Getenv("KMS_FAILPOINT_SERVERS_RUNNING")
+		if kmsTlsTestcase == "" {
+			mt.Skipf("Skipping test as KMS_FAILPOINT_SERVERS_RUNNING is not set")
+		}
+
 		setFailPoint := func(failure string, count int) error {
 			url := fmt.Sprintf("https://localhost:9003/set_failpoint/%s", failure)
 			var payloadBuf bytes.Buffer