package main

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func TestWSClientKeepaliveSendsPing(t *testing.T) {
	ping := make(chan struct{}, 1)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		conn.SetPingHandler(func(string) error {
			select {
			case ping <- struct{}{}:
			default:
			}
			return conn.WriteControl(websocket.PongMessage, nil, time.Now().Add(time.Second))
		})
		for {
			if _, _, err = conn.ReadMessage(); err != nil {
				return
			}
		}
	}))
	defer server.Close()

	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
	if err != nil {
		t.Fatal(err)
	}
	client := &WSClient{conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
	go client.readLoop()
	client.startKeepalive(10 * time.Millisecond)
	defer client.Close()

	select {
	case <-ping:
	case <-time.After(time.Second):
		t.Fatal("websocket keepalive ping was not received")
	}
}

func TestWSClientReadLoopClosesConnectionAfterDecodeError(t *testing.T) {
	serverSawClose := make(chan struct{}, 1)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		if err = conn.WriteMessage(websocket.TextMessage, []byte("{")); err != nil {
			return
		}
		if _, _, err = conn.ReadMessage(); err != nil {
			serverSawClose <- struct{}{}
		}
	}))
	defer server.Close()

	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
	if err != nil {
		t.Fatal(err)
	}
	client := &WSClient{conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
	go client.readLoop()

	select {
	case <-client.done:
	case <-time.After(time.Second):
		t.Fatal("read loop did not exit after malformed frame")
	}
	select {
	case <-serverSawClose:
	case <-time.After(time.Second):
		t.Fatal("read loop exited without closing the underlying connection")
	}
}