diff --git a/cmd/puri/main.go b/cmd/puri/main.go
index e572b4e..180b0e7 100644
--- a/cmd/puri/main.go
+++ b/cmd/puri/main.go
@@ -55,13 +55,13 @@ func main() {
 
 	switch mode {
 	case Param:
-		handleOutput(puri.ExtractParam(uri, *p))
+		handleOutput(puri.ExtractParam(handleInput(uri), *p))
 	case Scheme:
-		handleOutput(puri.ExtractScheme(uri))
+		handleOutput(puri.ExtractScheme(handleInput(uri)))
 	case Port:
-		handleOutput(puri.ExtractPort(uri))
+		handleOutput(puri.ExtractPort(handleInput(uri)))
 	case Host:
-		handleOutput(puri.ExtractHost(uri))
+		handleOutput(puri.ExtractHost(handleInput(uri)))
 	case Usage:
 
 		printUsage()
@@ -71,6 +71,14 @@ func main() {
 	}
 }
 
+func handleInput(uri string) url.URL {
+	parsed, err := url.Parse(uri)
+	if err != nil || len(uri) == 0 {
+		panic(fmt.Errorf("invalid uri: %v", uri))
+	}
+	return *parsed
+}
+
 func handleOutput(res *string, err error) {
 	if err == nil {
 		fmt.Println(*res)
diff --git a/puri.go b/puri.go
index f2680fc..2e2e4a7 100644
--- a/puri.go
+++ b/puri.go
@@ -8,78 +8,78 @@ import (
 
 const Version string = "v0.1.5"
 
-func parseURL(uri string) (*url.URL, error) {
-	parsed, err := url.Parse(uri)
-	if err != nil || len(uri) == 0 {
-		return nil, errors.New("invalid uri")
-	}
-	return parsed, nil
-}
+const colon = ":"
+const schemeSeparator = "://"
 
-func ExtractParam(uri string, param string) (*string, error) {
-	parsed, err := parseURL(uri)
-	if err != nil {
-		return nil, err
-	}
-	v := parsed.Query()
+// ExtractParam returns a pointer to a string containing the value of the specified parameter. Note,
+// returns empty strings if value not found so can be used for optional params
+func ExtractParam(uri url.URL, param string) (*string, error) {
+	v := uri.Query()
 	r := v.Get(param)
 	return &r, nil
 }
 
-func ExtractScheme(uri string) (*string, error) {
-	parsed, err := parseURL(uri)
-	if err != nil {
-		return nil, err
-	}
-
-	if parsed != nil && len(parsed.Scheme) == 0 {
-		return nil, errors.New("no scheme")
+// ExtractScheme extracts the scheme from a given URL.
+func ExtractScheme(uri url.URL) (*string, error) {
+	if len(uri.Scheme) == 0 {
+		return NoScheme()
 	}
 
-	return &parsed.Scheme, nil
+	return &uri.Scheme, nil
 }
 
-func ExtractHost(uri string) (*string, error) {
-	parsed, err := parseURL(uri)
-	if err != nil {
-		return nil, err
-	}
-
-	if parsed != nil && len(parsed.Host) == 0 {
-		if len(parsed.Path) > 0 {
-			return &parsed.Path, nil
+// ExtractHost extracts the host from the given URL. If the host is in the format "host:port", it splits the host
+// and returns only the host part.
+func ExtractHost(uri url.URL) (*string, error) {
+	if len(uri.Host) == 0 {
+		if len(uri.Path) > 0 {
+			return &uri.Path, nil
 		}
-		return nil, errors.New("no host")
+		return NoHost()
 	}
 
-	hp := strings.Split(parsed.Host, ":")
+	hp := strings.Split(uri.Host, colon)
 	return &hp[0], nil
 }
 
-func ExtractPort(uri string) (*string, error) {
-	parsed, err := parseURL(uri)
-	if err != nil {
-		return nil, err
-	}
-
-	if parsed != nil && len(parsed.Host) == 0 {
-		if len(parsed.Path) > 0 {
-			parsed.Host = parsed.Path
+// ExtractPort extracts the port from a given URL if present, otherwise returns an error.
+func ExtractPort(uri url.URL) (*string, error) {
+	if len(uri.Host) == 0 {
+		if len(uri.Path) > 0 {
+			uri.Host = uri.Path
 		}
-		if len(parsed.Scheme) > 0 {
-			parsed.Host = parsed.Scheme
+		if len(uri.Scheme) > 0 {
+			uri.Host = uri.Scheme
 		}
 	}
 
-	hp := strings.Split(parsed.Host, ":")
+	hp := strings.Split(uri.Host, colon)
 	if len(hp) != 2 {
-		if !strings.Contains(uri, "://") {
-			hp = strings.Split(uri, ":")
+		if !strings.Contains(uri.String(), schemeSeparator) {
+			hp = strings.Split(uri.String(), colon)
 		}
 		if len(hp) != 2 {
-			return nil, errors.New("no port")
+			return NoPort()
 		}
 	}
 
 	return &hp[1], nil
 }
+
+// NoPort returns an error indicating that there is no port associated with the URL. It is used
+// in the ExtractPort function to handle cases where the port cannot be extracted from the URL.
+func NoPort() (*string, error) {
+	p := ""
+	return &p, nil
+}
+
+// NoScheme returns an error indicating that no scheme was provided.
+// It returns a nil string pointer and an error with the message "no scheme".
+func NoScheme() (*string, error) {
+	return nil, errors.New("no scheme")
+}
+
+// NoHost returns an error indicating that there is no host available.
+func NoHost() (*string, error) {
+	return nil, errors.New("no host")
+}
diff --git a/puri_test.go b/puri_test.go
index 2870f28..87f15fc 100644
--- a/puri_test.go
+++ b/puri_test.go
@@ -1,6 +1,7 @@
 package puri
 
 import (
+	"net/url"
 	"testing"
 )
 
@@ -8,33 +9,41 @@ func p[T any](t T) *T {
 	return &t
 }
 
+func toTestURL(uri string) url.URL {
+	u, e := url.Parse(uri)
+	if e != nil {
+		panic("test setup failure")
+	}
+	return *u
+}
+
 func TestExtractParam(t *testing.T) {
 	tests := []struct {
 		name  string
-		uri   string
+		uri   url.URL
 		param string
 		want  *string
 	}{
 		{
 			name:  "case 1: valid url and param",
-			uri:   "http://example.com/?key=value",
+			uri:   toTestURL("http://example.com/?key=value"),
 			param: "key",
 			want:  p("value"),
 		},
 		{
 			name:  "case 2: valid url, param is missing",
-			uri:   "http://example.com/?key=value",
+			uri:   toTestURL("http://example.com/?key=value"),
 			param: "missing",
 		},
 		{
 			name:  "case 3: invalid url, valid param",
-			uri:   "http:/example.com?param1=value1",
+			uri:   toTestURL("http:/example.com?param1=value1"),
 			param: "param1",
 			want:  p("value1"),
 		},
 		{
 			name:  "case 4: empty url and param",
-			uri:   "",
+			uri:   toTestURL(""),
 			param: "",
 		},
 		// Add more cases as needed.
@@ -59,15 +68,15 @@ func TestExtractParam(t *testing.T) {
 func TestExtractScheme(t *testing.T) {
 	tests := []struct {
 		name       string
-		uri        string
+		uri        url.URL
 		wantScheme *string
 		wantError  bool
 	}{
-		{"ftp", "ftp://example.com", p("ftp"), false},
-		{"http", "http://example.com", p("http"), false},
-		{"https", "https://example.com", p("https"), false},
-		{"none", "empty uri", nil, true},
-		{"bad", "invalid uri", nil, true},
+		{"ftp", toTestURL("ftp://example.com"), p("ftp"), false},
+		{"http", toTestURL("http://example.com"), p("http"), false},
+		{"https", toTestURL("https://example.com"), p("https"), false},
+		{"none", toTestURL("empty uri"), nil, true},
+		{"bad", toTestURL("invalid uri"), nil, true},
 	}
 
 	for _, tc := range tests {
@@ -88,15 +97,15 @@ func TestExtractScheme(t *testing.T) {
 func TestExtractHost(t *testing.T) {
 	tests := []struct {
 		name      string
-		uri       string
+		uri       url.URL
 		wantHost  *string
 		wantError bool
 	}{
-		{"ftp host", "ftp://example.com", p("example.com"), false},
-		{"http host with port", "http://example.com:8080", p("example.com"), false},
-		{"http host with port and path", "http://example.com:8080/blah/blah?k=v", p("example.com"), false},
-		{"simple", "example.com", p("example.com"), false},
-		{"simpler", "host", p("host"), false},
+		{"ftp host", toTestURL("ftp://example.com"), p("example.com"), false},
+		{"http host with port", toTestURL("http://example.com:8080"), p("example.com"), false},
+		{"http host with port and path", toTestURL("http://example.com:8080/blah/blah?k=v"), p("example.com"), false},
+		{"simple", toTestURL("example.com"), p("example.com"), false},
+		{"simpler", toTestURL("host"), p("host"), false},
 	}
 
 	for _, tc := range tests {
@@ -116,15 +125,15 @@ func TestExtractHost(t *testing.T) {
 func TestExtractPort(t *testing.T) {
 	tests := []struct {
 		name      string
-		uri       string
+		uri       url.URL
 		wantPort  *string
 		wantError bool
 	}{
-		{"ftp host no port", "ftp://example.com", nil, true},
-		{"http host with port", "http://example.com:8080", p("8080"), false},
-		{"http host with port and path", "http://example.com:8080/blah/blah?k=v", p("8080"), false},
-		{"simple", "example.com:80", p("80"), false},
-		{"simpler", "host", nil, true},
+		{"ftp host no port", toTestURL("ftp://example.com"), p(""), false},
+		{"http host with port", toTestURL("http://example.com:8080"), p("8080"), false},
+		{"http host with port and path", toTestURL("http://example.com:8080/blah/blah?k=v"), p("8080"), false},
+		{"simple", toTestURL("example.com:80"), p("80"), false},
+		{"simpler", toTestURL("host"), p(""), false},
 	}
 
 	for _, tc := range tests {