update unified integration

Author: qingyang-hu

internal/integration/mtest/mongotest.go

@@ -797,12 +797,6 @@ func verifyRunOnBlockConstraint(rob RunOnBlock) error {
 		return err
 	}
 
-	// TODO(GODRIVER-3486): Once auto encryption is supported by the unified test
-	// format,this check should be removed.
-	if rob.CSFLEEnabled() && rob.CSFLE.Options != nil {
-		return fmt.Errorf("Auto encryption required (GODRIVER-3486)")
-	}
-
 	if rob.CSFLEEnabled() && !IsCSFLEEnabled() {
 		return fmt.Errorf("runOnBlock requires CSFLE to be enabled. Build with the cse tag to enable")
 	} else if !rob.CSFLEEnabled() && IsCSFLEEnabled() {

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,65 @@ func getURIForClient(opts *entityOptions) string {
 	}
 }
 
+func createAutoEncryptionOptions(opts bson.Raw) (*options.AutoEncryptionOptions, error) {
+	aeo := options.AutoEncryption()
+	var kvnsFound bool
+	elems, err := opts.Elements()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, elem := range elems {
+		name := elem.Key()
+		opt := elem.Value()
+
+		switch name {
+		case "kmsProviders":
+			providers := make(map[string]map[string]any)
+			elems, err := opt.Document().Elements()
+			if err != nil {
+				return nil, err
+			}
+			for _, elem := range elems {
+				provider := elem.Key()
+				providerOpt := elem.Value()
+				switch provider {
+				case "aws":
+					providers["aws"] = map[string]any{
+						"accessKeyId":     awsAccessKeyID,
+						"secretAccessKey": awsSecretAccessKey,
+					}
+				case "local":
+					_, key := providerOpt.Document().Lookup("key").Binary()
+					providers["local"] = map[string]any{
+						"key": key,
+					}
+				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"`

internal/spectest/skip.go

@@ -392,19 +392,6 @@ var skipTests = map[string][]string{
 		"TestClientSideEncryptionSpec/timeoutMS.json/timeoutMS_applied_to_listCollections_to_get_collection_schema",
 	},
 
-	// TODO(GODRIVER-3486): Support auto encryption in unified tests.
-	"Support auto encryption in unified tests (GODRIVER-3486)": {
-		"TestUnifiedSpec/unified-test-format/tests/valid-pass/poc-queryable-encryption.json/insert,_replace,_and_find_with_queryable_encryption",
-	},
-
-	// TODO(DRIVERS-3106): Support auto encryption in unified tests.
-	"Support auto encryption in unified tests (DRIVERS-3106)": {
-		"TestUnifiedSpec/client-side-encryption/tests/unified/localSchema.json/A_local_schema_should_override",
-		"TestUnifiedSpec/client-side-encryption/tests/unified/localSchema.json/A_local_schema_with_no_encryption_is_an_error",
-		"TestUnifiedSpec/client-side-encryption/tests/unified/fle2v2-BypassQueryAnalysis.json/BypassQueryAnalysis_decrypts",
-		"TestUnifiedSpec/client-side-encryption/tests/unified/fle2v2-EncryptedFields-vs-EncryptedFieldsMap.json/encryptedFieldsMap_is_preferred_over_remote_encryptedFields",
-	},
-
 	// TODO(GODRIVER-3076): CSFLE/QE Support for more than 1 KMS provider per
 	// type.
 	"Support multiple KMS providers per type (GODRIVER-3076)": {