update test case

Author: qingyang-hu

internal/integration/client_side_encryption_prose_test.go

@@ -3177,11 +3177,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 						return cred, nil
 					},
 				})
-			clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
-			assert.NoErrorf(mt, err, "error on NewClientEncryption: %v", err)
-
-			dkOpts := options.DataKey()
-			_, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
+			_, err = mongo.NewClientEncryption(keyVaultClient, ceo)
 			assert.Error(mt, err, "expected an error")
 		})
 		mt.Run("Case 2: ClientEncryption with credentialProviders works", func(mt *mtest.T) {
@@ -3209,7 +3205,10 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
 			assert.NoErrorf(mt, err, "error on NewClientEncryption: %v", err)
 
-			dkOpts := options.DataKey()
+			dkOpts := options.DataKey().SetMasterKey(bson.D{
+				{"region", "us-east-1"},
+				{"key", "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0"},
+			})
 			_, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
 			assert.NoErrorf(mt, err, "unexpected error %v", err)
 			assert.Equal(mt, 1, calledCount, "expected credential provider to be called once")
@@ -3254,35 +3253,32 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			keyVaultClient, err := mongo.Connect(opts)
 			assert.NoErrorf(mt, err, "error on Connect: %v", err)
 
+			var calledCount int
 			ceo := options.ClientEncryption().
 				SetKeyVaultNamespace("keyvault.datakeys").
 				SetKmsProviders(map[string]map[string]any{
-					"aws": {
-						"accessKeyId":     awsAccessKeyID,
-						"secretAccessKey": awsSecretAccessKey,
-					},
+					"aws": map[string]any{},
 				}).
 				SetCredentialProviders(map[string]options.CredentialsProvider{
 					"aws": func(ctx context.Context) (options.Credentials, error) {
-						var cred options.Credentials
-						provider := credproviders.NewEnvProvider()
-						c, err := provider.Retrieve(ctx)
-						if err != nil {
-							return cred, err
-						}
-						cred.AccessKeyID = c.AccessKeyID
-						cred.SecretAccessKey = c.SecretAccessKey
-						cred.SessionToken = c.SessionToken
-						cred.ExpirationCallback = provider.IsExpired
-						return cred, nil
+						calledCount++
+						return options.Credentials{
+							AccessKeyID:        awsAccessKeyID,
+							SecretAccessKey:    awsSecretAccessKey,
+							ExpirationCallback: func() bool { return false },
+						}, nil
 					},
 				})
 			clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
 			assert.NoErrorf(mt, err, "error on NewClientEncryption: %v", err)
 
-			dkOpts := options.DataKey()
+			dkOpts := options.DataKey().SetMasterKey(bson.D{
+				{"region", "us-east-1"},
+				{"key", "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0"},
+			})
 			_, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
 			assert.NoErrorf(mt, err, "unexpected error %v", err)
+			assert.Equal(mt, 1, calledCount, "expected credential provider to be called once")
 		})
 	})
 }

internal/test/aws/aws_test.go

@@ -13,6 +13,7 @@ import (
 	"testing"
 
 	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/internal/credproviders"
 	"go.mongodb.org/mongo-driver/v2/internal/require"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
@@ -46,15 +47,20 @@ func TestAWSCustomCredentialProviders(t *testing.T) {
 	}
 
 	var calledCount int
+	provider := credproviders.NewEnvProvider()
 	awsCredential := options.Credential{
-		AuthMechanism: "MONGODB-AWS",
-		AwsCredentialsProvider: func(_ context.Context) (options.Credentials, error) {
+		AwsCredentialsProvider: func(ctx context.Context) (options.Credentials, error) {
 			calledCount++
-			return options.Credentials{
-				AccessKeyID:        os.Getenv("AWS_ACCESS_KEY_ID"),
-				SecretAccessKey:    os.Getenv("AWS_SECRET_ACCESS_KEY"),
-				ExpirationCallback: func() bool { return false },
-			}, nil
+			var creds options.Credentials
+			value, err := provider.Retrieve(ctx)
+			if err != nil {
+				return creds, err
+			}
+			creds.AccessKeyID = value.AccessKeyID
+			creds.SecretAccessKey = value.SecretAccessKey
+			creds.SessionToken = value.SessionToken
+			creds.ExpirationCallback = provider.IsExpired
+			return creds, nil
 		},
 	}
 	client, err := mongo.Connect(options.Client().ApplyURI(uri).SetAuth(awsCredential))

x/mongo/driver/mongocrypt/mongocrypt.go

@@ -65,12 +65,16 @@ func NewMongoCrypt(opts *options.MongoCryptOptions) (*MongoCrypt, error) {
 	if needsKmsProvider(opts.KmsProviders, "gcp") {
 		kmsProviders["gcp"] = creds.NewGCPCredentialProvider(httpClient)
 	}
+	provider, ok := opts.CredentialProviders["aws"]
 	if needsKmsProvider(opts.KmsProviders, "aws") {
 		var providers []credentials.Provider
-		if provider, ok := opts.CredentialProviders["aws"]; ok {
+		if ok {
 			providers = append(providers, provider)
 		}
 		kmsProviders["aws"] = creds.NewAWSCredentialProvider(httpClient, providers...)
+	} else if ok {
+		return nil, fmt.Errorf("can only provide a custom AWS credential provider " +
+			"when the state machine is configured for automatic AWS credential fetching")
 	}
 	if needsKmsProvider(opts.KmsProviders, "azure") {
 		kmsProviders["azure"] = creds.NewAzureCredentialProvider(httpClient)