update unified integration

Author: qingyang-hu

internal/integration/unified/client_entity.go

@@ -9,6 +9,7 @@ package unified
 import (
 	"context"
 	"fmt"
+	"os"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -32,11 +33,16 @@ import (
 // exceed the default truncation length.
 const defaultMaxDocumentLen = 10_000
 
-// Security-sensitive commands that should be ignored in command monitoring by default.
-var securitySensitiveCommands = []string{
-	"authenticate", "saslStart", "saslContinue", "getnonce",
-	"createUser", "updateUser", "copydbgetnonce", "copydbsaslstart", "copydb",
-}
+var (
+	// Security-sensitive commands that should be ignored in command monitoring by default.
+	securitySensitiveCommands = []string{
+		"authenticate", "saslStart", "saslContinue", "getnonce",
+		"createUser", "updateUser", "copydbgetnonce", "copydbsaslstart", "copydb",
+	}
+
+	awsAccessKeyID     = os.Getenv("FLE_AWS_KEY")
+	awsSecretAccessKey = os.Getenv("FLE_AWS_SECRET")
+)
 
 // clientEntity is a wrapper for a mongo.Client object that also holds additional information required during test
 // execution.
@@ -217,6 +223,13 @@ func newClientEntity(ctx context.Context, em *EntityMap, entityOptions *entityOp
 	} else {
 		integtest.AddTestServerAPIVersion(clientOpts)
 	}
+	if entityOptions.AutoEncryptOpts != nil {
+		aeo, err := createAutoEncryptionOptions(entityOptions.AutoEncryptOpts)
+		if err != nil {
+			return nil, fmt.Errorf("error parsing auto encryption options: %w", err)
+		}
+		clientOpts.SetAutoEncryptionOptions(aeo)
+	}
 	for _, cmd := range entityOptions.IgnoredCommands {
 		entity.ignoredCommands[cmd] = struct{}{}
 	}
@@ -251,6 +264,53 @@ func getURIForClient(opts *entityOptions) string {
 	}
 }
 
+func createAutoEncryptionOptions(opts bson.Raw) (*options.AutoEncryptionOptions, error) {
+	aeo := options.AutoEncryption()
+	var kvnsFound bool
+	elems, _ := opts.Elements()
+
+	for _, elem := range elems {
+		name := elem.Key()
+		opt := elem.Value()
+
+		switch name {
+		case "kmsProviders":
+			providers := make(map[string]map[string]any)
+			elems, _ := opt.Document().Elements()
+			for _, elem := range elems {
+				provider := elem.Key()
+				switch provider {
+				case "aws":
+					providers["aws"] = map[string]any{
+						"accessKeyId":     awsAccessKeyID,
+						"secretAccessKey": awsSecretAccessKey,
+					}
+				default:
+					return nil, fmt.Errorf("unrecognized KMS provider: %v", provider)
+				}
+			}
+			aeo.SetKmsProviders(providers)
+		case "schemaMap":
+			var schemaMap map[string]any
+			err := bson.Unmarshal(opt.Document(), &schemaMap)
+			if err != nil {
+				return nil, err
+			}
+			aeo.SetSchemaMap(schemaMap)
+		case "keyVaultNamespace":
+			kvnsFound = true
+			aeo.SetKeyVaultNamespace(opt.StringValue())
+		default:
+			return nil, fmt.Errorf("unrecognized option: %v", name)
+		}
+	}
+	if !kvnsFound {
+		aeo.SetKeyVaultNamespace("keyvault.datakeys")
+	}
+
+	return aeo, nil
+}
+
 // disconnect disconnects the client associated with this entity. It is an
 // idempotent operation, unlike the mongo client's disconnect method. This
 // property will help avoid unnecessary errors when calling disconnect on a

internal/integration/unified/entity.go

@@ -52,6 +52,7 @@ type entityOptions struct {
 	ID string `bson:"id"`
 
 	// Options for client entities.
+	AutoEncryptOpts          bson.Raw                      `bson:"autoEncryptOpts"`
 	URIOptions               bson.M                        `bson:"uriOptions"`
 	UseMultipleMongoses      *bool                         `bson:"useMultipleMongoses"`
 	ObserveEvents            []string                      `bson:"observeEvents"`