From 8aaf5b8f98d69f3931a1e4ec1e9316a0c05a2b78 Mon Sep 17 00:00:00 2001
From: "P. Reis" <patrickpereirareal1@gmail.com>
Date: Thu, 4 Jul 2024 12:06:24 -0300
Subject: [PATCH] Append 'ws://' prefix if url starts with localhost (#132)

* test: refactor NormalizeURL tests

* feat(NormalizeURL): add 'ws://' prefix if url starts with 'localhost'

* test(NormalizeURL): add 'ws://' prefix if url is localhost
---
 normalize.go      |  4 ++++
 normalize_test.go | 60 +++++++++++++++++++++++++----------------------
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/normalize.go b/normalize.go
index 168da12..6e9333b 100644
--- a/normalize.go
+++ b/normalize.go
@@ -15,6 +15,10 @@ func NormalizeURL(u string) string {
 	u = strings.TrimSpace(u)
 	u = strings.ToLower(u)
 
+	if strings.HasPrefix(u, "localhost") == true {
+		u = "ws://" + u
+	}
+
 	if !strings.HasPrefix(u, "http") && !strings.HasPrefix(u, "ws") {
 		u = "wss://" + u
 	}
diff --git a/normalize_test.go b/normalize_test.go
index d6ee479..6b8d2ad 100644
--- a/normalize_test.go
+++ b/normalize_test.go
@@ -1,32 +1,36 @@
 package nostr
 
-import "fmt"
+import (
+	"testing"
+)
 
-func ExampleNormalizeURL() {
-	fmt.Println(NormalizeURL(""))
-	fmt.Println(NormalizeURL("wss://x.com/y"))
-	fmt.Println(NormalizeURL("wss://x.com/y/"))
-	fmt.Println(NormalizeURL("http://x.com/y"))
-	fmt.Println(NormalizeURL(NormalizeURL("http://x.com/y")))
-	fmt.Println(NormalizeURL("wss://x.com"))
-	fmt.Println(NormalizeURL("wss://x.com/"))
-	fmt.Println(NormalizeURL(NormalizeURL(NormalizeURL("wss://x.com/"))))
-	fmt.Println(NormalizeURL("x.com"))
-	fmt.Println(NormalizeURL("x.com/"))
-	fmt.Println(NormalizeURL("x.com////"))
-	fmt.Println(NormalizeURL("x.com/?x=23"))
-
-	// Output:
-	//
-	// wss://x.com/y
-	// wss://x.com/y
-	// ws://x.com/y
-	// ws://x.com/y
-	// wss://x.com
-	// wss://x.com
-	// wss://x.com
-	// wss://x.com
-	// wss://x.com
-	// wss://x.com
-	// wss://x.com?x=23
+type urlTest struct {
+	url, expected string
+}
+
+var urlTests = []urlTest{
+	{"", ""},
+	{"wss://x.com/y", "wss://x.com/y"},
+	{"wss://x.com/y/", "wss://x.com/y"},
+	{"http://x.com/y", "ws://x.com/y"},
+	{NormalizeURL("http://x.com/y"), "ws://x.com/y"},
+	{NormalizeURL("wss://x.com"), "wss://x.com"},
+	{NormalizeURL("wss://x.com/"), "wss://x.com"},
+	{NormalizeURL(NormalizeURL(NormalizeURL("wss://x.com/"))), "wss://x.com"},
+	{"wss://x.com", "wss://x.com"},
+	{"wss://x.com/", "wss://x.com"},
+	{"x.com////", "wss://x.com"},
+	{"x.com/?x=23", "wss://x.com?x=23"},
+	{"localhost:4036", "ws://localhost:4036"},
+	{"localhost:4036/relay", "ws://localhost:4036/relay"},
+	{NormalizeURL("localhost:4036/relay"), "ws://localhost:4036/relay"},
+}
+
+func TestNormalizeURL(t *testing.T) {
+
+	for _, test := range urlTests {
+		if output := NormalizeURL(test.url); output != test.expected {
+			t.Errorf("Output '%s' not equal to expected '%s'", output, test.expected)
+		}
+	}
 }