diff --git a/app.go b/app.go
index e2d6e82..53706a9 100644
--- a/app.go
+++ b/app.go
@@ -7,7 +7,7 @@ import (
 	"time"
 
 	"github.com/saml-dev/gome-assistant/internal/http"
-	"github.com/saml-dev/gome-assistant/internal/setup"
+	ws "github.com/saml-dev/gome-assistant/internal/websocket"
 	"nhooyr.io/websocket"
 )
 
@@ -30,7 +30,7 @@ you can use to register schedules and listeners.
 */
 func App(connString string) app {
 	token := os.Getenv("AUTH_TOKEN")
-	conn, ctx, ctxCancel := setup.SetupConnection(connString)
+	conn, ctx, ctxCancel := ws.SetupConnection(connString)
 
 	httpClient := http.NewHttpClient(connString, token)
 
diff --git a/internal/services/light.go b/internal/services/light.go
index 0bac36a..2eac12e 100644
--- a/internal/services/light.go
+++ b/internal/services/light.go
@@ -4,7 +4,7 @@ import (
 	"context"
 
 	"github.com/saml-dev/gome-assistant/internal/http"
-	"github.com/saml-dev/gome-assistant/internal/setup"
+	ws "github.com/saml-dev/gome-assistant/internal/websocket"
 	"nhooyr.io/websocket"
 )
 
@@ -28,12 +28,12 @@ type LightRequest struct {
 
 func (l Light) TurnOn(entityId string) {
 	req := newLightOnRequest(entityId)
-	setup.WriteMessage(req, l.conn, l.ctx)
+	ws.WriteMessage(req, l.conn, l.ctx)
 }
 
 func (l Light) TurnOff(entityId string) {
 	req := newLightOffRequest(entityId)
-	setup.WriteMessage(req, l.conn, l.ctx)
+	ws.WriteMessage(req, l.conn, l.ctx)
 }
 
 /* Internal */
diff --git a/internal/setup/setup.go b/internal/setup/setup.go
deleted file mode 100644
index a9d0297..0000000
--- a/internal/setup/setup.go
+++ /dev/null
@@ -1,104 +0,0 @@
-package setup
-
-import (
-	"context"
-	"encoding/json"
-	"errors"
-	"fmt"
-	"log"
-	"os"
-	"time"
-
-	"nhooyr.io/websocket"
-)
-
-type AuthMessage struct {
-	MsgType     string `json:"type"`
-	AccessToken string `json:"access_token"`
-}
-
-func WriteMessage[T any](msg T, conn *websocket.Conn, ctx context.Context) error {
-	msgJson, err := json.Marshal(msg)
-	// fmt.Println(string(msgJson))
-	if err != nil {
-		return err
-	}
-
-	err = conn.Write(ctx, websocket.MessageText, msgJson)
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
-func ReadMessage(conn *websocket.Conn, ctx context.Context) (string, error) {
-	_, msg, err := conn.Read(ctx)
-	if err != nil {
-		return "", err
-	}
-	return string(msg), nil
-}
-
-func SetupConnection(connString string) (*websocket.Conn, context.Context, context.CancelFunc) {
-	ctx, ctxCancel := context.WithTimeout(context.Background(), time.Second*3)
-
-	// Init websocket connection
-	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://%s/api/websocket", connString), nil)
-	if err != nil {
-		ctxCancel()
-		log.Fatalf("ERROR: Failed to connect to websocket at ws://%s/api/websocket. Check IP address and port\n", connString)
-	}
-
-	// Read auth_required message
-	_, err = ReadMessage(conn, ctx)
-	if err != nil {
-		ctxCancel()
-		log.Fatalln("Unknown error creating websocket client")
-	}
-
-	// Send auth message
-	err = SendAuthMessage(conn, ctx)
-	if err != nil {
-		ctxCancel()
-		log.Fatalln("Unknown error creating websocket client")
-	}
-
-	// Verify auth message
-	err = VerifyAuthResponse(conn, ctx)
-	if err != nil {
-		ctxCancel()
-		log.Fatalln("ERROR: Auth token is invalid. Please double check it or create a new token in your Home Assistant profile")
-	}
-
-	return conn, ctx, ctxCancel
-}
-
-func SendAuthMessage(conn *websocket.Conn, ctx context.Context) error {
-	token := os.Getenv("AUTH_TOKEN")
-	err := WriteMessage(AuthMessage{MsgType: "auth", AccessToken: token}, conn, ctx)
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
-type authResponse struct {
-	MsgType string `json:"type"`
-	Message string `json:"message"`
-}
-
-func VerifyAuthResponse(conn *websocket.Conn, ctx context.Context) error {
-	_, msg, err := conn.Read(ctx)
-	if err != nil {
-		return err
-	}
-
-	var authResp authResponse
-	json.Unmarshal(msg, &authResp)
-	if authResp.MsgType != "auth_ok" {
-		return errors.New("invalid auth token")
-	}
-
-	return nil
-}
diff --git a/internal/websocket/websocket.go b/internal/websocket/websocket.go
index ca8837c..984cf22 100644
--- a/internal/websocket/websocket.go
+++ b/internal/websocket/websocket.go
@@ -1,7 +1,108 @@
-// websocket is used to interact with the Home Assistant
+// Package websocket is used to interact with the Home Assistant
 // websocket API. All HA interaction is done via websocket
 // except for cases explicitly called out in http package
 // documentation.
 package websocket
 
-// TODO: move existing websocket code here
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"os"
+	"time"
+
+	"nhooyr.io/websocket"
+)
+
+type AuthMessage struct {
+	MsgType     string `json:"type"`
+	AccessToken string `json:"access_token"`
+}
+
+func WriteMessage[T any](msg T, conn *websocket.Conn, ctx context.Context) error {
+	msgJson, err := json.Marshal(msg)
+	// fmt.Println(string(msgJson))
+	if err != nil {
+		return err
+	}
+
+	err = conn.Write(ctx, websocket.MessageText, msgJson)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func ReadMessage(conn *websocket.Conn, ctx context.Context) (string, error) {
+	_, msg, err := conn.Read(ctx)
+	if err != nil {
+		return "", err
+	}
+	return string(msg), nil
+}
+
+func SetupConnection(connString string) (*websocket.Conn, context.Context, context.CancelFunc) {
+	ctx, ctxCancel := context.WithTimeout(context.Background(), time.Second*3)
+
+	// Init websocket connection
+	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://%s/api/websocket", connString), nil)
+	if err != nil {
+		ctxCancel()
+		log.Fatalf("ERROR: Failed to connect to websocket at ws://%s/api/websocket. Check IP address and port\n", connString)
+	}
+
+	// Read auth_required message
+	_, err = ReadMessage(conn, ctx)
+	if err != nil {
+		ctxCancel()
+		log.Fatalln("Unknown error creating websocket client")
+	}
+
+	// Send auth message
+	err = SendAuthMessage(conn, ctx)
+	if err != nil {
+		ctxCancel()
+		log.Fatalln("Unknown error creating websocket client")
+	}
+
+	// Verify auth message
+	err = VerifyAuthResponse(conn, ctx)
+	if err != nil {
+		ctxCancel()
+		log.Fatalln("ERROR: Auth token is invalid. Please double check it or create a new token in your Home Assistant profile")
+	}
+
+	return conn, ctx, ctxCancel
+}
+
+func SendAuthMessage(conn *websocket.Conn, ctx context.Context) error {
+	token := os.Getenv("AUTH_TOKEN")
+	err := WriteMessage(AuthMessage{MsgType: "auth", AccessToken: token}, conn, ctx)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+type authResponse struct {
+	MsgType string `json:"type"`
+	Message string `json:"message"`
+}
+
+func VerifyAuthResponse(conn *websocket.Conn, ctx context.Context) error {
+	_, msg, err := conn.Read(ctx)
+	if err != nil {
+		return err
+	}
+
+	var authResp authResponse
+	json.Unmarshal(msg, &authResp)
+	if authResp.MsgType != "auth_ok" {
+		return errors.New("invalid auth token")
+	}
+
+	return nil
+}