package web_console

import (
	"bytes"
	"io"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func TestWSBridge_DoesNotCloseUntilBothStreamsDone(t *testing.T) {
	clientConn, serverConn, cleanup := newWebSocketPair(t)
	defer cleanup()

	stderrR, stderrW := io.Pipe()
	defer func() {
		_ = stderrR.Close()
		_ = stderrW.Close()
	}()

	sshCli := &sshClient{
		conn:     serverConn,
		sessOut:  bytes.NewReader(nil),
		sessErr:  stderrR,
		closeSig: make(chan struct{}, 1),
	}

	stdoutDone := make(chan struct{})
	go func() {
		_ = sshCli.wsWrite()
		close(stdoutDone)
	}()
	go func() {
		_ = sshCli.wsWriteStderr()
	}()
	go func() {
		<-sshCli.closeSig
		_ = serverConn.Close()
	}()

	select {
	case <-stdoutDone:
	case <-time.After(2 * time.Second):
		t.Fatalf("stdout writer did not exit")
	}

	if _, err := stderrW.Write([]byte("err\n")); err != nil {
		t.Fatalf("stderrW.Write: %v", err)
	}
	_ = stderrW.Close()

	_ = clientConn.SetReadDeadline(time.Now().Add(2 * time.Second))
	msgType, payload, err := clientConn.ReadMessage()
	if err != nil {
		t.Fatalf("clientConn.ReadMessage: %v", err)
	}
	if msgType != websocket.BinaryMessage {
		t.Fatalf("expected BinaryMessage, got %d", msgType)
	}
	if string(payload) != "err\n" {
		t.Fatalf("unexpected payload: %q", string(payload))
	}
}