| 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 | } |
| 74 |
Members