From 1cdf1a6ae72f32b0ede18c1515122590abfadb02 Mon Sep 17 00:00:00 2001
From: zhangruizhi <the2ndhitler@163.com>
Date: Fri, 1 Sep 2023 12:40:24 +0800
Subject: [PATCH] extract whole `buildURI` logic and write a test

---
 main.go      | 29 +++++++++-------
 main_test.go | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+), 12 deletions(-)

diff --git a/main.go b/main.go
index 966c88850..2dc302199 100644
--- a/main.go
+++ b/main.go
@@ -91,6 +91,22 @@ func main() {
 	e.Run()
 }
 
+func buildURI(uri string, user string, password string) string {
+	// IF user@pass not contained in uri AND custom user and pass supplied in arguments
+	// DO concat a new uri with user and pass arguments value
+	if !strings.Contains(uri, "@") && user != "" && password != "" {
+		// trim mongodb:// prefix to handle user and pass logic
+		uri = strings.TrimPrefix(uri, "mongodb://")
+
+		// log.Debugf("add user and pass to the uri")
+		uri = fmt.Sprintf("%s:%s@%s", user, password, uri)
+
+		// add back mongodb://
+		uri = "mongodb://" + uri
+	}
+	return uri
+}
+
 func buildExporter(opts GlobalFlags) *exporter.Exporter {
 	log := logrus.New()
 
@@ -105,18 +121,7 @@ func buildExporter(opts GlobalFlags) *exporter.Exporter {
 
 	log.Debugf("Compatible mode: %v", opts.CompatibleMode)
 
-	// IF user@pass not contained in uri AND custom user and pass supplied in arguments
-	// DO concat a new uri with user and pass arguments value
-	if !strings.Contains(opts.URI, "@") && opts.User != "" && opts.Password != "" {
-		// trim mongodb:// prefix to handle user and pass logic
-		opts.URI = strings.TrimPrefix(opts.URI, "mongodb://")
-
-		log.Debugf("add user and pass to the uri")
-		opts.URI = fmt.Sprintf("%s:%s@%s", opts.User, opts.Password, opts.URI)
-
-		// add back mongodb://
-		opts.URI = "mongodb://" + opts.URI
-	}
+	opts.URI = buildURI(opts.URI, opts.User, opts.Password)
 
 	log.Debugf("Connection URI: %s", opts.URI)
 
diff --git a/main_test.go b/main_test.go
index e11f468bd..f4db34142 100644
--- a/main_test.go
+++ b/main_test.go
@@ -38,3 +38,98 @@ func TestBuildExporter(t *testing.T) {
 
 	buildExporter(opts)
 }
+
+func TestBuildURI(t *testing.T) {
+	const newUser = "xxx"
+	const newPass = "yyy"
+
+	const originalBareURI = "127.0.0.1"
+	const originalAuthURI = "usr:pwd@127.0.0.1"
+
+	const originalPrefixBareURI = "mongodb://127.0.0.1"
+	const originalPrefixAuthURI = "mongodb://usr:pwd@127.0.0.1"
+	const changedPrefixAuthURI = "mongodb://xxx:yyy@127.0.0.1"
+
+	var newUri string
+
+	t.Log("\nuri with prefix and auth, and auth supplied in opt.User/Password")
+	newUri = buildURI(originalPrefixAuthURI, newUser, newPass)
+	t.Logf("Origin: %s", originalPrefixAuthURI)
+	t.Logf("Expect: %s", originalPrefixAuthURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != originalPrefixAuthURI {
+		t.Fail()
+	}
+	newUri = ""
+
+	t.Log("\nuri with prefix and auth, no auth supplied in opt.User/Password")
+	newUri = buildURI(originalPrefixAuthURI, "", "")
+	t.Logf("Origin: %s", originalPrefixAuthURI)
+	t.Logf("Expect: %s", originalPrefixAuthURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != originalPrefixAuthURI {
+		t.Fail()
+	}
+	newUri = ""
+
+	t.Log("\nuri with no prefix and auth, and auth supplied in opt.User/Password")
+	newUri = buildURI(originalAuthURI, newUser, newPass)
+	t.Logf("Origin: %s", originalAuthURI)
+	t.Logf("Expect: %s", originalAuthURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != originalAuthURI {
+		t.Fail()
+	}
+	newUri = ""
+
+	t.Log("\nuri with no prefix and auth, no auth supplied in opt.User/Password")
+	newUri = buildURI(originalAuthURI, "", "")
+	t.Logf("Origin: %s", originalAuthURI)
+	t.Logf("Expect: %s", originalAuthURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != originalAuthURI {
+		t.Fail()
+	}
+	newUri = ""
+
+	t.Log("\nuri with prefix and no auth, and auth supplied in opt.User/Password")
+	newUri = buildURI(originalPrefixBareURI, newUser, newPass)
+	t.Logf("Origin: %s", originalPrefixBareURI)
+	t.Logf("Expect: %s", changedPrefixAuthURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != changedPrefixAuthURI {
+		t.Fail()
+	}
+	newUri = ""
+
+	t.Log("\nuri with prefix and no auth, no auth supplied in opt.User/Password")
+	newUri = buildURI(originalPrefixBareURI, "", "")
+	t.Logf("Origin: %s", originalPrefixBareURI)
+	t.Logf("Expect: %s", originalPrefixBareURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != originalPrefixBareURI {
+		t.Fail()
+	}
+	newUri = ""
+
+	t.Log("\nuri with no prefix and no auth, and auth supplied in opt.User/Password")
+	newUri = buildURI(originalBareURI, newUser, newPass)
+	t.Logf("Origin: %s", originalBareURI)
+	t.Logf("Expect: %s", changedPrefixAuthURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != changedPrefixAuthURI {
+		t.Fail()
+	}
+	newUri = ""
+
+	t.Log("\nuri with no prefix and no auth, no auth supplied in opt.User/Password")
+	newUri = buildURI(originalBareURI, "", "")
+	t.Logf("Origin: %s", originalBareURI)
+	t.Logf("Expect: %s", originalBareURI)
+	t.Logf("Result: %s", newUri)
+	if newUri != originalBareURI {
+		t.Fail()
+	}
+	newUri = ""
+
+}