1 files changed, 5 insertions, 4 deletions

diff --git a/api/cmd/main_test.go b/api/cmd/main_test.go
index a6c8381..049cf6e 100644
--- a/api/cmd/main_test.go
+++ b/api/cmd/main_test.go
@@ -6,6 +6,7 @@ import (
         "net/http"
         "net/http/httptest"
         "testing"
+        "water/api/internal/router"
 
         "github.com/spf13/viper"
         "github.com/stretchr/testify/assert"
@@ -26,7 +27,7 @@ func getTestUserCredentials() (string, string) {
 }
 
 func TestAuthRoute(t *testing.T) {
-        router := setupRouter()
+        r := router.SetupRouter()
         
         username, password := getTestUserCredentials()
 
@@ -39,7 +40,7 @@ func TestAuthRoute(t *testing.T) {
                 t.Fatalf("Failed to create request: %v", err)
         }
         req.SetBasicAuth(username, password)
-        router.ServeHTTP(w, req)
+        r.ServeHTTP(w, req)
 
         assert.Equal(t, http.StatusOK, w.Code, "response should return a 200 code")
 
@@ -56,12 +57,12 @@ func TestAuthRoute(t *testing.T) {
 }
 
 func TestAuthRouteFailure(t *testing.T) {
-        router := setupRouter()
+        r := router.SetupRouter()
 
         w := httptest.NewRecorder()
         req, _ := http.NewRequest("POST", "/api/v1/auth", nil)
         req.SetBasicAuth("asdf", "asdf")
-        router.ServeHTTP(w, req)
+        r.ServeHTTP(w, req)
 
         assert.Equal(t, http.StatusUnauthorized, w.Code, "should return a 401 code")
 }