Clean up code

Author: prestonvasquez

internal/cmd/testatlas/atlas_test.go

@@ -8,18 +8,20 @@ package main
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/base64"
 	"errors"
 	"fmt"
 	"net/url"
 	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/internal/handshake"
-	"go.mongodb.org/mongo-driver/v2/internal/require"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
@@ -29,91 +31,91 @@ func TestAtlas(t *testing.T) {
 		name        string
 		envVar      string
 		certKeyFile string
-		wantErr     string
+		wantErrCode string
 	}{
 		{
 			name:        "Atlas with TLS",
 			envVar:      "ATLAS_REPL",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with TLS and shared cluster",
 			envVar:      "ATLAS_SHRD",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with free tier",
 			envVar:      "ATLAS_FREE",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with TLS 1.1",
 			envVar:      "ATLAS_TLS11",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with TLS 1.2",
 			envVar:      "ATLAS_TLS12",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with serverless",
 			envVar:      "ATLAS_SERVERLESS",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with srv file on replica set",
 			envVar:      "ATLAS_SRV_REPL",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with srv file on shared cluster",
 			envVar:      "ATLAS_SRV_SHRD",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with srv file on free tier",
 			envVar:      "ATLAS_SRV_FREE",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with srv file on TLS 1.1",
 			envVar:      "ATLAS_SRV_TLS11",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with srv file on TLS 1.2",
 			envVar:      "ATLAS_SRV_TLS12",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with srv file on serverless",
 			envVar:      "ATLAS_SRV_SERVERLESS",
 			certKeyFile: "",
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with X509 Dev",
 			envVar:      "ATLAS_X509_DEV",
 			certKeyFile: createAtlasX509DevCertKeyFile(t),
-			wantErr:     "",
+			wantErrCode: "",
 		},
 		{
 			name:        "Atlas with X509 Dev no user",
 			envVar:      "ATLAS_X509_DEV",
 			certKeyFile: createAtlasX509DevCertKeyFileNoUser(t),
-			wantErr:     "(UserNotFound) Could not find user",
+			wantErrCode: "11", // UserNotFound
 		},
 	}
 
@@ -133,21 +135,34 @@ func TestAtlas(t *testing.T) {
 
 			// Run basic connectivity test.
 			err := runTest(context.Background(), clientOpts)
-			if tc.wantErr != "" {
-				assert.ErrorContains(t, err, tc.wantErr, "error running test with TLS")
+			if tc.wantErrCode != "" {
+				var cmdErr *mongo.CommandError
+				if errors.As(err, &cmdErr) {
+					assert.Equal(t, cmdErr.Code, tc.wantErrCode)
+				} else {
+					t.Fatalf("expected error to be a command error, got: %v", err)
+				}
 
 				return
 			}
 			require.NoError(t, err, "error running test with TLS")
 
-			tlsConfigSkipVerify := clientOpts.TLSConfig
-			tlsConfigSkipVerify.InsecureSkipVerify = true
+			orig := clientOpts.TLSConfig
+			if orig == nil {
+				orig = &tls.Config{}
+			}
+
+			insecure := orig.Clone()
+			insecure.InsecureSkipVerify = true
 
 			// Run the connectivity test with InsecureSkipVerify to ensure SNI is done
 			// correctly even if verification is disabled.
-			clientOpts.SetTLSConfig(tlsConfigSkipVerify)
+			insecureClientOpts := options.Client().
+				ApplyURI(uri).
+				SetServerSelectionTimeout(1 * time.Second).
+				SetTLSConfig(insecure)
 
-			err = runTest(context.Background(), clientOpts)
+			err = runTest(context.Background(), insecureClientOpts)
 			require.NoError(t, err, "error running test with tlsInsecure")
 		})
 	}
@@ -180,7 +195,10 @@ func runTest(ctx context.Context, clientOpts *options.ClientOptions) error {
 func createAtlasX509DevCertKeyFile(t *testing.T) string {
 	t.Helper()
 
-	certBytes, err := base64.StdEncoding.DecodeString(os.Getenv("ATLAS_X509_DEV_CERT_BASE64"))
+	b64 := os.Getenv("ATLAS_X509_DEV_CERT_BASE64")
+	assert.NotEmpty(t, b64, "Environment variable ATLAS_X509_DEV_CERT_BASE64 is not set")
+
+	certBytes, err := base64.StdEncoding.DecodeString(b64)
 	require.NoError(t, err, "failed to decode ATLAS_X509_DEV_CERT_BASE64")
 
 	certFilePath := t.TempDir() + "/atlas_x509_dev_cert.pem"
@@ -194,7 +212,10 @@ func createAtlasX509DevCertKeyFile(t *testing.T) string {
 func createAtlasX509DevCertKeyFileNoUser(t *testing.T) string {
 	t.Helper()
 
-	keyBytes, err := base64.StdEncoding.DecodeString(os.Getenv("ATLAS_X509_DEV_CERT_NOUSER_BASE64"))
+	b64 := os.Getenv("ATLAS_X509_DEV_CERT_NOUSER_BASE64")
+	assert.NotEmpty(t, b64, "Environment variable ATLAS_X509_DEV_CERT_NOUSER_BASE64 is not set")
+
+	keyBytes, err := base64.StdEncoding.DecodeString(b64)
 	require.NoError(t, err, "failed to decode ATLAS_X509_DEV_CERT_NOUSER_BASE64")
 
 	keyFilePath := t.TempDir() + "/atlas_x509_dev_cert_no_user.pem"
@@ -212,7 +233,7 @@ func addTLSCertKeyFile(t *testing.T, certKeyFile, uri string) string {
 	require.NoError(t, err, "failed to parse uri")
 
 	q := u.Query()
-	q.Set("tlsCertificateKeyFile", certKeyFile)
+	q.Set("tlsCertificateKeyFile", filepath.ToSlash(certKeyFile))
 
 	u.RawQuery = q.Encode()