diff --git a/cmd/server/main.go b/cmd/server/main.go
index 57caf5e..58199b2 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -7,6 +7,7 @@ import (
 	"log"
 	"net/http"
 	"os/signal"
+	"strings"
 	"sync"
 	"syscall"
 	"time"
@@ -71,11 +72,20 @@ func main() {
 	r.HandleFunc("/api/settings", controller.SettingsListController(appState, client, ctx)).Methods("GET")
 	r.HandleFunc("/api/settings", controller.SettingsEditController(settingsWriter, appState, client, ctx)).Methods("POST")
 
-	// r.HandleFunc("/api/beacons/ws", serveWs(appState, ctx))
+	wsHandler := http.HandlerFunc(serveWs(appState, ctx))
+	restApiHandler := handlers.CORS(originsOk, headersOk, methodsOk)(r)
+	mainHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if strings.HasPrefix(r.URL.Path, "/api/beacons/ws") {
+			wsHandler.ServeHTTP(w, r)
+			return
+		}
+
+		restApiHandler.ServeHTTP(w, r)
+	})
 
 	server := http.Server{
 		Addr:    cfg.HTTPAddr,
-		Handler: handlers.CORS(originsOk, headersOk, methodsOk)(r),
+		Handler: mainHandler,
 	}
 
 	go server.ListenAndServe()
@@ -119,9 +129,9 @@ func serveWs(appstate *appcontext.AppState, ctx context.Context) http.HandlerFun
 			}
 			return
 		}
-		wg.Add(1)
+		wg.Add(2)
 		go writer(ws, appstate, ctx)
-		reader(ws)
+		go reader(ws, ctx)
 	}
 }
 
@@ -161,17 +171,25 @@ func writer(ws *websocket.Conn, appstate *appcontext.AppState, ctx context.Conte
 	}
 }
 
-func reader(ws *websocket.Conn) {
+func reader(ws *websocket.Conn, ctx context.Context) {
 	defer func() {
 		ws.Close()
+		wg.Done()
 	}()
 	ws.SetReadLimit(512)
 	ws.SetReadDeadline(time.Now().Add((60 * 9) / 10 * time.Second))
 	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add((60 * 9) / 10 * time.Second)); return nil })
 	for {
-		_, _, err := ws.ReadMessage()
-		if err != nil {
-			break
+		fmt.Println("this is called")
+		select {
+		case <-ctx.Done():
+			fmt.Println("closing websocket reader")
+			return
+		default:
+			_, _, err := ws.ReadMessage()
+			if err != nil {
+				return
+			}
 		}
 	}
 }
diff --git a/scripts/testAPI.sh b/scripts/testAPI.sh
index 3ee97f6..0611905 100755
--- a/scripts/testAPI.sh
+++ b/scripts/testAPI.sh
@@ -8,6 +8,8 @@ curl -s -X POST $URL \
   -d '{"Beacon_id":"'"$BEACON_ID"'","Name":"Beacon1","tx_power":-59,"rssi":-70}'
 echo -e "\n"
 
+sleep 1
+
 curl -X GET $URL
 
 sleep 1
\ No newline at end of file