github.com/cockroachdb/tools@v0.0.0-20230222021103-a6d27438930d/playground/socket/socket_test.go (about)

     1  // Copyright 2015 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package socket
     6  
     7  import (
     8  	"testing"
     9  	"time"
    10  )
    11  
    12  func TestBuffer(t *testing.T) {
    13  	afterChan := make(chan time.Time)
    14  	ch := make(chan *Message)
    15  	go func() {
    16  		ch <- &Message{Kind: "err", Body: "a"}
    17  		ch <- &Message{Kind: "err", Body: "b"}
    18  		ch <- &Message{Kind: "out", Body: "1"}
    19  		ch <- &Message{Kind: "out", Body: "2"}
    20  		afterChan <- time.Time{} // value itself doesn't matter
    21  		ch <- &Message{Kind: "out", Body: "3"}
    22  		ch <- &Message{Kind: "out", Body: "4"}
    23  		close(ch)
    24  	}()
    25  
    26  	var ms []*Message
    27  	timeAfter := func(d time.Duration) <-chan time.Time {
    28  		return afterChan
    29  	}
    30  	for m := range buffer(ch, timeAfter) {
    31  		ms = append(ms, m)
    32  	}
    33  	if len(ms) != 3 {
    34  		t.Fatalf("got %v messages, want 3", len(ms))
    35  	}
    36  	if g, w := ms[0].Body, "ab"; g != w {
    37  		t.Errorf("message 0 body = %q, want %q", g, w)
    38  	}
    39  	if g, w := ms[1].Body, "12"; g != w {
    40  		t.Errorf("message 1 body = %q, want %q", g, w)
    41  	}
    42  	if g, w := ms[2].Body, "34"; g != w {
    43  		t.Errorf("message 2 body = %q, want %q", g, w)
    44  	}
    45  }
    46  
    47  type killRecorder chan struct{}
    48  
    49  func (k killRecorder) Kill() { close(k) }
    50  
    51  func TestLimiter(t *testing.T) {
    52  	ch := make(chan *Message)
    53  	go func() {
    54  		var m Message
    55  		for i := 0; i < msgLimit+10; i++ {
    56  			ch <- &m
    57  		}
    58  		ch <- &Message{Kind: "end"}
    59  	}()
    60  
    61  	kr := make(killRecorder)
    62  	n := 0
    63  	for m := range limiter(ch, kr) {
    64  		n++
    65  		if n > msgLimit && m.Kind != "end" {
    66  			t.Errorf("received non-end message after limit")
    67  		}
    68  	}
    69  	if n != msgLimit+1 {
    70  		t.Errorf("received %v messages, want %v", n, msgLimit+1)
    71  	}
    72  	<-kr
    73  }