Source file
src/iter/pull_test.go
Documentation: iter
1
2
3
4
5
6
7 package iter
8
9 import (
10 "fmt"
11 "runtime"
12 "testing"
13 )
14
15 func count(n int) Seq[int] {
16 return func(yield func(int) bool) {
17 for i := range n {
18 if !yield(i) {
19 break
20 }
21 }
22 }
23 }
24
25 func squares(n int) Seq2[int, int64] {
26 return func(yield func(int, int64) bool) {
27 for i := range n {
28 if !yield(i, int64(i)*int64(i)) {
29 break
30 }
31 }
32 }
33 }
34
35 func TestPull(t *testing.T) {
36
37 for end := 0; end <= 3; end++ {
38 t.Run(fmt.Sprint(end), func(t *testing.T) {
39 ng := runtime.NumGoroutine()
40 wantNG := func(want int) {
41 if xg := runtime.NumGoroutine() - ng; xg != want {
42 t.Helper()
43 t.Errorf("have %d extra goroutines, want %d", xg, want)
44 }
45 }
46 wantNG(0)
47 next, stop := Pull(count(3))
48 wantNG(1)
49 for i := range end {
50 v, ok := next()
51 if v != i || ok != true {
52 t.Fatalf("next() = %d, %v, want %d, %v", v, ok, i, true)
53 }
54 wantNG(1)
55 }
56 wantNG(1)
57 if end < 3 {
58 stop()
59 wantNG(0)
60 }
61 for range 2 {
62 v, ok := next()
63 if v != 0 || ok != false {
64 t.Fatalf("next() = %d, %v, want %d, %v", v, ok, 0, false)
65 }
66 wantNG(0)
67 }
68 wantNG(0)
69
70 stop()
71 stop()
72 stop()
73 wantNG(0)
74 })
75 }
76 }
77
78 func TestPull2(t *testing.T) {
79 for end := 0; end <= 3; end++ {
80 t.Run(fmt.Sprint(end), func(t *testing.T) {
81 ng := runtime.NumGoroutine()
82 wantNG := func(want int) {
83 if xg := runtime.NumGoroutine() - ng; xg != want {
84 t.Helper()
85 t.Errorf("have %d extra goroutines, want %d", xg, want)
86 }
87 }
88 wantNG(0)
89 next, stop := Pull2(squares(3))
90 wantNG(1)
91 for i := range end {
92 k, v, ok := next()
93 if k != i || v != int64(i*i) || ok != true {
94 t.Fatalf("next() = %d, %d, %v, want %d, %d, %v", k, v, ok, i, i*i, true)
95 }
96 wantNG(1)
97 }
98 wantNG(1)
99 if end < 3 {
100 stop()
101 wantNG(0)
102 }
103 for range 2 {
104 k, v, ok := next()
105 if v != 0 || ok != false {
106 t.Fatalf("next() = %d, %d, %v, want %d, %d, %v", k, v, ok, 0, 0, false)
107 }
108 wantNG(0)
109 }
110 wantNG(0)
111
112 stop()
113 stop()
114 stop()
115 wantNG(0)
116 })
117 }
118 }
119
View as plain text