<commit_before>package juggler

import (
	"bytes"
	"encoding/json"
	"io"
	"sync"
	"testing"
	"time"

	"github.com/PuerkitoBio/exp/juggler/internal/wstest"
	"github.com/PuerkitoBio/exp/juggler/msg"
	"github.com/gorilla/websocket"
	"github.com/pborman/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestClient(t *testing.T) {
	var buf bytes.Buffer
	done := make(chan bool, 1)
	srv := wstest.StartRecordingServer(t, done, &buf)
	defer srv.Close()

	// the only received message should be EXP
	var (
		mu         sync.Mutex
		cnt        int
		expForUUID uuid.UUID
		wg         sync.WaitGroup
	)
	h := MsgHandlerFunc(func(m msg.Msg) {
		defer wg.Done()

		mu.Lock()
		cnt++
		if assert.Equal(t, msg.ExpMsg, m.Type(), "Expects EXP message") {
			expForUUID = m.(*msg.Exp).Payload.For
		}
		mu.Unlock()
	})

	cli, err := Dial(&websocket.Dialer{}, srv.URL, nil, h)
	require.NoError(t, err, "Dial")
	cli.CallTimeout = time.Millisecond
	cli.LogFunc = (&debugLog{t: t}).Printf

	// call
	wg.Add(1)
	type expected struct {
		uid uuid.UUID
		mt  msg.MessageType
	}
	var expectedResults []expected
	callUUID, err := cli.Call("a", "call", 0)
	require.NoError(t, err, "Call")
	expectedResults = append(expectedResults, expected{callUUID, msg.CallMsg})

	uid, err := cli.Pub("b", "pub")
	require.NoError(t, err, "Pub")
	expectedResults = append(expectedResults, expected{uid, msg.PubMsg})

	uid, err = cli.Sub("c", false)
	require.NoError(t, err, "Sub")
	expectedResults = append(expectedResults, expected{uid, msg.SubMsg})

	uid, err = cli.Unsb("d", true)
	require.NoError(t, err, "Unsb")
	expectedResults = append(expectedResults, expected{uid, msg.UnsbMsg})

	// wait for any pending handlers
	wg.Wait()

	cli.Close()
	<-done

	mu.Lock()
	defer mu.Unlock()
	assert.Equal(t, 1, cnt, "Expected calls to Handler")
	assert.Equal(t, callUUID, expForUUID, "Expired message should be for the Call message")

	// read the messages received by the server
	var p json.RawMessage
	r := bytes.NewReader(buf.Bytes())
	dec := json.NewDecoder(r)
	for i, exp := range expectedResults {
		require.NoError(t, dec.Decode(&p), "Decode %d", i)
		m, err := msg.Unmarshal(bytes.NewReader(p))<commit_msg>juggler: fix race when running TestClient with TestUpgrade<commit_after>package juggler

import (
	"bytes"
	"encoding/json"
	"io"
	"sync"
	"testing"
	"time"

	"github.com/PuerkitoBio/exp/juggler/internal/wstest"
	"github.com/PuerkitoBio/exp/juggler/msg"
	"github.com/gorilla/websocket"
	"github.com/pborman/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestClient(t *testing.T) {
	var buf bytes.Buffer
	done := make(chan bool, 1)
	srv := wstest.StartRecordingServer(t, done, &buf)
	defer srv.Close()

	// the only received message should be EXP
	var (
		mu         sync.Mutex
		cnt        int
		expForUUID uuid.UUID
		wg         sync.WaitGroup
	)
	h := MsgHandlerFunc(func(m msg.Msg) {
		defer wg.Done()

		mu.Lock()
		cnt++
		if assert.Equal(t, msg.ExpMsg, m.Type(), "Expects EXP message") {
			expForUUID = m.(*msg.Exp).Payload.For
		}
		mu.Unlock()
	})

	closed := make(chan bool)
	dbgClientClosed = func(c *Client) { closed <- true }
	defer func() { dbgClientClosed = nil }()

	cli, err := Dial(&websocket.Dialer{}, srv.URL, nil, h)
	require.NoError(t, err, "Dial")
	cli.CallTimeout = time.Millisecond
	cli.LogFunc = (&debugLog{t: t}).Printf

	// call
	wg.Add(1)
	type expected struct {
		uid uuid.UUID
		mt  msg.MessageType
	}
	var expectedResults []expected
	callUUID, err := cli.Call("a", "call", 0)
	require.NoError(t, err, "Call")
	expectedResults = append(expectedResults, expected{callUUID, msg.CallMsg})

	uid, err := cli.Pub("b", "pub")
	require.NoError(t, err, "Pub")
	expectedResults = append(expectedResults, expected{uid, msg.PubMsg})

	uid, err = cli.Sub("c", false)
	require.NoError(t, err, "Sub")
	expectedResults = append(expectedResults, expected{uid, msg.SubMsg})

	uid, err = cli.Unsb("d", true)
	require.NoError(t, err, "Unsb")
	expectedResults = append(expectedResults, expected{uid, msg.UnsbMsg})

	// wait for any pending handlers
	wg.Wait()

	cli.Close()
	<-done
	<-closed

	mu.Lock()
	defer mu.Unlock()
	assert.Equal(t, 1, cnt, "Expected calls to Handler")
	assert.Equal(t, callUUID, expForUUID, "Expired message should be for the Call message")

	// read the messages received by the server
	var p json.RawMessage
	r := bytes.NewReader(buf.Bytes())
	dec := json.NewDecoder(r)
	for i, exp := range expectedResults {
		require.NoError(t, dec.Decode(&p), "Decode %d", i)
		m, err := msg.Unmarshal(bytes.NewReader(p))
{'fraction_non_alphanumeric': 0.15351163296368775, 'fraction_numerical': 0.0017395085888236572, 'mean_word_length': 4.043859649122807, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4461015', 'n_tokens_mistral': 1807, 'n_tokens_neox': 1741, 'n_words': 401}