aboutsummaryrefslogtreecommitdiff
path: root/pipeline
diff options
context:
space:
mode:
Diffstat (limited to 'pipeline')
-rw-r--r--pipeline/batch.go64
-rw-r--r--pipeline/batch_test.go106
-rw-r--r--pipeline/example_test.go53
-rw-r--r--pipeline/options.go34
-rw-r--r--pipeline/pipeline.go21
-rw-r--r--pipeline/pipeline_bench_test.go55
-rw-r--r--pipeline/pipeline_test.go255
-rw-r--r--pipeline/stage.go310
-rw-r--r--pipeline/stage_test.go308
9 files changed, 1206 insertions, 0 deletions
diff --git a/pipeline/batch.go b/pipeline/batch.go
new file mode 100644
index 0000000..74b08f0
--- /dev/null
+++ b/pipeline/batch.go
@@ -0,0 +1,64 @@
1package pipeline
2
3import (
4 "context"
5 "time"
6)
7
8// Batch collects items into slices of up to size, flushing when full or when timeout elapses.
9func Batch[T any](ctx context.Context, in <-chan T, size int, timeout time.Duration) <-chan []T {
10 if in == nil {
11 panic("pipeline: Batch input channel must not be nil")
12 }
13 if size < 1 {
14 panic("pipeline: Batch size must be at least 1")
15 }
16
17 out := make(chan []T)
18
19 go func() {
20 defer close(out)
21
22 buf := make([]T, 0, size)
23 ticker := time.NewTicker(timeout)
24 defer ticker.Stop()
25
26 flush := func() bool {
27 if len(buf) == 0 {
28 return true
29 }
30 batch := buf
31 buf = make([]T, 0, size)
32 select {
33 case out <- batch:
34 return true
35 case <-ctx.Done():
36 return false
37 }
38 }
39
40 for {
41 select {
42 case v, ok := <-in:
43 if !ok {
44 flush()
45 return
46 }
47 buf = append(buf, v)
48 if len(buf) >= size {
49 if !flush() {
50 return
51 }
52 }
53 case <-ticker.C:
54 if !flush() {
55 return
56 }
57 case <-ctx.Done():
58 return
59 }
60 }
61 }()
62
63 return out
64}
diff --git a/pipeline/batch_test.go b/pipeline/batch_test.go
new file mode 100644
index 0000000..4d79ddc
--- /dev/null
+++ b/pipeline/batch_test.go
@@ -0,0 +1,106 @@
1package pipeline
2
3import (
4 "context"
5 "testing"
6 "time"
7)
8
9func TestBatch_Full(t *testing.T) {
10 t.Parallel()
11
12 input := make(chan int, 10)
13 for i := 1; i <= 10; i++ {
14 input <- i
15 }
16 close(input)
17
18 out := Batch(t.Context(), input, 3, time.Second)
19
20 var batches [][]int
21 for b := range out {
22 batches = append(batches, b)
23 }
24
25 if len(batches) != 4 {
26 t.Fatalf("expected 4 batches, got %d: %v", len(batches), batches)
27 }
28 if len(batches[0]) != 3 {
29 t.Fatalf("expected first batch size 3, got %d", len(batches[0]))
30 }
31 if len(batches[3]) != 1 {
32 t.Fatalf("expected last batch size 1, got %d", len(batches[3]))
33 }
34}
35
36func TestBatch_TimeoutFlush(t *testing.T) {
37 t.Parallel()
38
39 input := make(chan int)
40
41 out := Batch(t.Context(), input, 100, 50*time.Millisecond)
42
43 input <- 1
44 input <- 2
45
46 select {
47 case batch := <-out:
48 if len(batch) != 2 {
49 t.Fatalf("expected batch of 2 from timeout flush, got %d", len(batch))
50 }
51 case <-time.After(time.Second):
52 t.Fatal("timeout waiting for batch flush")
53 }
54
55 close(input)
56 for range out {
57 }
58}
59
60func TestBatch_ContextCancel(t *testing.T) {
61 t.Parallel()
62 ctx, cancel := context.WithCancel(t.Context())
63
64 input := make(chan int)
65 out := Batch(ctx, input, 100, time.Second)
66
67 cancel()
68
69 done := make(chan struct{})
70 go func() {
71 for range out {
72 }
73 close(done)
74 }()
75
76 select {
77 case <-done:
78 case <-time.After(time.Second):
79 t.Fatal("timeout waiting for shutdown")
80 }
81}
82
83func TestBatch_NilInput(t *testing.T) {
84 t.Parallel()
85
86 defer func() {
87 if r := recover(); r == nil {
88 t.Error("expected panic for nil input")
89 }
90 }()
91
92 Batch[int](t.Context(), nil, 10, time.Second)
93}
94
95func TestBatch_PanicOnZeroSize(t *testing.T) {
96 t.Parallel()
97
98 defer func() {
99 if r := recover(); r == nil {
100 t.Error("expected panic for batch size < 1")
101 }
102 }()
103
104 input := make(chan int)
105 Batch(t.Context(), input, 0, time.Second)
106}
diff --git a/pipeline/example_test.go b/pipeline/example_test.go
new file mode 100644
index 0000000..e322529
--- /dev/null
+++ b/pipeline/example_test.go
@@ -0,0 +1,53 @@
1package pipeline_test
2
3import (
4 "context"
5 "fmt"
6 "sort"
7
8 "git.sangtran.io.vn/quacker.git/pipeline"
9)
10
11func ExampleRun() {
12 input := make(chan int, 5)
13 for i := 1; i <= 5; i++ {
14 input <- i
15 }
16 close(input)
17
18 out := pipeline.Run(context.Background(), input,
19 pipeline.Filter(func(v int) bool { return v%2 == 0 }),
20 pipeline.Map(1, func(_ context.Context, v int) (int, bool) {
21 return v * 10, true
22 }),
23 )
24
25 var results []int
26 for v := range out {
27 results = append(results, v)
28 }
29 sort.Ints(results)
30 fmt.Println(results)
31 // Output: [20 40]
32}
33
34func ExampleSink() {
35 input := make(chan string, 3)
36 input <- "hello"
37 input <- "world"
38 input <- "!"
39 close(input)
40
41 var collected []string
42 out := pipeline.Run(context.Background(), input,
43 pipeline.Sink(1, func(_ context.Context, v string) {
44 collected = append(collected, v)
45 }),
46 )
47
48 for range out {
49 }
50
51 fmt.Println(collected)
52 // Output: [hello world !]
53}
diff --git a/pipeline/options.go b/pipeline/options.go
new file mode 100644
index 0000000..c8c616e
--- /dev/null
+++ b/pipeline/options.go
@@ -0,0 +1,34 @@
1package pipeline
2
3// StageObserver receives callbacks for stage lifecycle events.
4type StageObserver interface {
5 OnItem()
6 OnDrop()
7 OnDone()
8}
9
10// StageOption configures a stage's behavior.
11type StageOption func(*stageConfig)
12
13type stageConfig struct {
14 observer StageObserver
15 bufferSize int
16}
17
18// WithObserver attaches an observer to a stage.
19func WithObserver(obs StageObserver) StageOption {
20 return func(cfg *stageConfig) { cfg.observer = obs }
21}
22
23// WithBufferSize sets the output channel buffer size for a stage.
24func WithBufferSize(size int) StageOption {
25 return func(cfg *stageConfig) { cfg.bufferSize = size }
26}
27
28func applyOpts(opts []StageOption) stageConfig {
29 var cfg stageConfig
30 for _, opt := range opts {
31 opt(&cfg)
32 }
33 return cfg
34}
diff --git a/pipeline/pipeline.go b/pipeline/pipeline.go
new file mode 100644
index 0000000..170a04e
--- /dev/null
+++ b/pipeline/pipeline.go
@@ -0,0 +1,21 @@
1// Package pipeline provides composable, type-safe stream processing with backpressure.
2package pipeline
3
4import "context"
5
6// Stage transforms an input channel into an output channel.
7// The stage owns the output channel and closes it when the input
8// is exhausted or the context is cancelled.
9type Stage[T any] func(ctx context.Context, in <-chan T) <-chan T
10
11// Run connects stages into a pipeline and returns the final output channel.
12func Run[T any](ctx context.Context, input <-chan T, stages ...Stage[T]) <-chan T {
13 if input == nil {
14 panic("pipeline: input channel must not be nil")
15 }
16 ch := input
17 for _, stage := range stages {
18 ch = stage(ctx, ch)
19 }
20 return ch
21}
diff --git a/pipeline/pipeline_bench_test.go b/pipeline/pipeline_bench_test.go
new file mode 100644
index 0000000..f4f81a0
--- /dev/null
+++ b/pipeline/pipeline_bench_test.go
@@ -0,0 +1,55 @@
1package pipeline
2
3import (
4 "context"
5 "testing"
6)
7
8func BenchmarkMap(b *testing.B) {
9 b.ReportAllocs()
10
11 input := make(chan int, 256)
12
13 out := Run(b.Context(), input,
14 Map(4, func(_ context.Context, v int) (int, bool) {
15 return v * 2, true
16 }),
17 )
18
19 go func() {
20 for range out {
21 }
22 }()
23
24 for b.Loop() {
25 input <- 1
26 }
27
28 close(input)
29}
30
31func BenchmarkPipeline_Composition(b *testing.B) {
32 b.ReportAllocs()
33
34 input := make(chan int, 256)
35
36 out := Run(b.Context(), input,
37 Filter(func(v int) bool { return v > 0 }),
38 Map(4, func(_ context.Context, v int) (int, bool) {
39 return v * 2, true
40 }),
41 Buffer[int](64),
42 Sink(2, func(_ context.Context, v int) {}),
43 )
44
45 go func() {
46 for range out {
47 }
48 }()
49
50 for b.Loop() {
51 input <- 1
52 }
53
54 close(input)
55}
diff --git a/pipeline/pipeline_test.go b/pipeline/pipeline_test.go
new file mode 100644
index 0000000..9593356
--- /dev/null
+++ b/pipeline/pipeline_test.go
@@ -0,0 +1,255 @@
1package pipeline
2
3import (
4 "context"
5 "fmt"
6 "sync/atomic"
7 "testing"
8 "time"
9)
10
11func TestRun_Passthrough(t *testing.T) {
12 t.Parallel()
13
14 input := make(chan int, 3)
15 input <- 1
16 input <- 2
17 input <- 3
18 close(input)
19
20 out := Run(t.Context(), input)
21
22 var got []int
23 for v := range out {
24 got = append(got, v)
25 }
26
27 if len(got) != 3 {
28 t.Fatalf("expected 3 items, got %d", len(got))
29 }
30}
31
32func TestRun_NilInput(t *testing.T) {
33 t.Parallel()
34
35 defer func() {
36 if r := recover(); r == nil {
37 t.Error("expected panic for nil input")
38 }
39 }()
40
41 Run[int](t.Context(), nil)
42}
43
44func TestRun_Composition(t *testing.T) {
45 t.Parallel()
46
47 input := make(chan int, 10)
48 for i := 1; i <= 10; i++ {
49 input <- i
50 }
51 close(input)
52
53 out := Run(t.Context(), input,
54 Filter(func(v int) bool { return v%2 == 0 }),
55 Map(1, func(_ context.Context, v int) (int, bool) {
56 return v * 2, true
57 }),
58 )
59
60 var got []int
61 for v := range out {
62 got = append(got, v)
63 }
64
65 if len(got) != 5 {
66 t.Fatalf("expected 5 items, got %d: %v", len(got), got)
67 }
68
69 sum := 0
70 for _, v := range got {
71 sum += v
72 }
73 if sum != 60 {
74 t.Fatalf("expected sum 60, got %d", sum)
75 }
76}
77
78func TestRun_Backpressure(t *testing.T) {
79 t.Parallel()
80
81 input := make(chan int)
82 consumed := make(chan int, 10)
83
84 out := Run(t.Context(), input,
85 Map(1, func(_ context.Context, v int) (int, bool) {
86 consumed <- v
87 return v, true
88 }),
89 )
90
91 input <- 42
92 v := <-consumed
93 if v != 42 {
94 t.Fatalf("expected 42, got %d", v)
95 }
96
97 select {
98 case input <- 99:
99 t.Fatal("producer was not blocked")
100 case <-time.After(50 * time.Millisecond):
101 }
102
103 go func() {
104 for range out {
105 }
106 }()
107 close(input)
108}
109
110func TestRun_GracefulDrain(t *testing.T) {
111 t.Parallel()
112
113 input := make(chan int, 5)
114 for i := 1; i <= 5; i++ {
115 input <- i
116 }
117 close(input)
118
119 var count atomic.Int32
120 out := Run(t.Context(), input,
121 Sink(2, func(_ context.Context, v int) {
122 count.Add(1)
123 }),
124 )
125
126 for range out {
127 }
128
129 if count.Load() != 5 {
130 t.Fatalf("expected all 5 items drained, got %d", count.Load())
131 }
132}
133
134func TestRun_Cancellation(t *testing.T) {
135 t.Parallel()
136 ctx, cancel := context.WithCancel(t.Context())
137
138 input := make(chan int)
139 out := Run(ctx, input,
140 Map(1, func(_ context.Context, v int) (int, bool) {
141 return v, true
142 }),
143 )
144
145 go func() {
146 select {
147 case input <- 1:
148 case <-ctx.Done():
149 }
150 }()
151
152 cancel()
153
154 done := make(chan struct{})
155 go func() {
156 for range out {
157 }
158 close(done)
159 }()
160
161 select {
162 case <-done:
163 case <-time.After(time.Second):
164 t.Fatal("pipeline did not shut down after cancellation")
165 }
166}
167
168func TestRun_ErrorPropagation(t *testing.T) {
169 t.Parallel()
170 ctx, cancel := context.WithCancelCause(t.Context())
171 defer cancel(nil)
172
173 input := make(chan int, 5)
174 for i := 1; i <= 5; i++ {
175 input <- i
176 }
177 close(input)
178
179 sentinel := fmt.Errorf("bad item")
180
181 out := Run(ctx, input,
182 Map(1, func(_ context.Context, v int) (int, bool) {
183 if v == 3 {
184 cancel(sentinel)
185 return v, false
186 }
187 return v * 2, true
188 }),
189 )
190
191 for range out {
192 }
193
194 if cause := context.Cause(ctx); cause != sentinel {
195 t.Fatalf("expected sentinel error, got: %v", cause)
196 }
197}
198
199func TestObserver(t *testing.T) {
200 t.Parallel()
201
202 var items, drops atomic.Int32
203 doneCh := make(chan struct{})
204
205 obs := &testObserver{
206 onItem: func() { items.Add(1) },
207 onDrop: func() { drops.Add(1) },
208 onDone: func() { close(doneCh) },
209 }
210
211 input := make(chan int, 6)
212 for i := 1; i <= 6; i++ {
213 input <- i
214 }
215 close(input)
216
217 out := Run(t.Context(), input,
218 Filter(func(v int) bool { return v%2 == 0 }, WithObserver(obs)),
219 )
220
221 for range out {
222 }
223 <-doneCh
224
225 if items.Load() != 3 {
226 t.Errorf("expected 3 items, got %d", items.Load())
227 }
228 if drops.Load() != 3 {
229 t.Errorf("expected 3 drops, got %d", drops.Load())
230 }
231}
232
233type testObserver struct {
234 onItem func()
235 onDrop func()
236 onDone func()
237}
238
239func (o *testObserver) OnItem() {
240 if o.onItem != nil {
241 o.onItem()
242 }
243}
244
245func (o *testObserver) OnDrop() {
246 if o.onDrop != nil {
247 o.onDrop()
248 }
249}
250
251func (o *testObserver) OnDone() {
252 if o.onDone != nil {
253 o.onDone()
254 }
255}
diff --git a/pipeline/stage.go b/pipeline/stage.go
new file mode 100644
index 0000000..e0716cf
--- /dev/null
+++ b/pipeline/stage.go
@@ -0,0 +1,310 @@
1package pipeline
2
3import (
4 "context"
5 "sync"
6)
7
8// Map applies fn to each item with bounded concurrency.
9// If keep is false, the item is dropped.
10func Map[T any](concurrency int, fn func(context.Context, T) (T, bool), opts ...StageOption) Stage[T] {
11 if concurrency < 1 {
12 panic("pipeline: Map concurrency must be at least 1")
13 }
14 cfg := applyOpts(opts)
15
16 return func(ctx context.Context, in <-chan T) <-chan T {
17 out := make(chan T, cfg.bufferSize)
18 var wg sync.WaitGroup
19
20 for range concurrency {
21 wg.Go(func() {
22 for {
23 select {
24 case v, ok := <-in:
25 if !ok {
26 return
27 }
28 result, keep := fn(ctx, v)
29 if !keep {
30 if cfg.observer != nil {
31 cfg.observer.OnDrop()
32 }
33 continue
34 }
35 if cfg.observer != nil {
36 cfg.observer.OnItem()
37 }
38 select {
39 case out <- result:
40 case <-ctx.Done():
41 return
42 }
43 case <-ctx.Done():
44 return
45 }
46 }
47 })
48 }
49
50 go func() {
51 wg.Wait()
52 close(out)
53 if cfg.observer != nil {
54 cfg.observer.OnDone()
55 }
56 }()
57
58 return out
59 }
60}
61
62// Filter forwards items for which fn returns true.
63func Filter[T any](fn func(T) bool, opts ...StageOption) Stage[T] {
64 cfg := applyOpts(opts)
65
66 return func(ctx context.Context, in <-chan T) <-chan T {
67 out := make(chan T, cfg.bufferSize)
68
69 go func() {
70 defer close(out)
71 defer func() {
72 if cfg.observer != nil {
73 cfg.observer.OnDone()
74 }
75 }()
76
77 for {
78 select {
79 case v, ok := <-in:
80 if !ok {
81 return
82 }
83 if !fn(v) {
84 if cfg.observer != nil {
85 cfg.observer.OnDrop()
86 }
87 continue
88 }
89 if cfg.observer != nil {
90 cfg.observer.OnItem()
91 }
92 select {
93 case out <- v:
94 case <-ctx.Done():
95 return
96 }
97 case <-ctx.Done():
98 return
99 }
100 }
101 }()
102
103 return out
104 }
105}
106
107// ForEach calls fn on each item and forwards it unchanged.
108func ForEach[T any](fn func(T), opts ...StageOption) Stage[T] {
109 cfg := applyOpts(opts)
110
111 return func(ctx context.Context, in <-chan T) <-chan T {
112 out := make(chan T, cfg.bufferSize)
113
114 go func() {
115 defer close(out)
116 defer func() {
117 if cfg.observer != nil {
118 cfg.observer.OnDone()
119 }
120 }()
121
122 for {
123 select {
124 case v, ok := <-in:
125 if !ok {
126 return
127 }
128 fn(v)
129 if cfg.observer != nil {
130 cfg.observer.OnItem()
131 }
132 select {
133 case out <- v:
134 case <-ctx.Done():
135 return
136 }
137 case <-ctx.Done():
138 return
139 }
140 }
141 }()
142
143 return out
144 }
145}
146
147// Buffer inserts a buffered channel of the specified capacity between stages.
148func Buffer[T any](size int) Stage[T] {
149 if size < 1 {
150 panic("pipeline: Buffer size must be at least 1")
151 }
152
153 return func(ctx context.Context, in <-chan T) <-chan T {
154 out := make(chan T, size)
155
156 go func() {
157 defer close(out)
158 for {
159 select {
160 case v, ok := <-in:
161 if !ok {
162 return
163 }
164 select {
165 case out <- v:
166 case <-ctx.Done():
167 return
168 }
169 case <-ctx.Done():
170 return
171 }
172 }
173 }()
174
175 return out
176 }
177}
178
179// Tee sends each item to a side channel while forwarding it to the next stage.
180func Tee[T any](side chan<- T, opts ...StageOption) Stage[T] {
181 if side == nil {
182 panic("pipeline: Tee side channel must not be nil")
183 }
184 cfg := applyOpts(opts)
185
186 return func(ctx context.Context, in <-chan T) <-chan T {
187 out := make(chan T, cfg.bufferSize)
188
189 go func() {
190 defer close(out)
191 defer func() {
192 if cfg.observer != nil {
193 cfg.observer.OnDone()
194 }
195 }()
196
197 for {
198 select {
199 case v, ok := <-in:
200 if !ok {
201 return
202 }
203 if cfg.observer != nil {
204 cfg.observer.OnItem()
205 }
206 select {
207 case side <- v:
208 case <-ctx.Done():
209 return
210 }
211 select {
212 case out <- v:
213 case <-ctx.Done():
214 return
215 }
216 case <-ctx.Done():
217 return
218 }
219 }
220 }()
221
222 return out
223 }
224}
225
226// Take forwards up to n items, then closes the output channel.
227func Take[T any](n int, opts ...StageOption) Stage[T] {
228 if n < 0 {
229 panic("pipeline: Take count must not be negative")
230 }
231 cfg := applyOpts(opts)
232
233 return func(ctx context.Context, in <-chan T) <-chan T {
234 out := make(chan T, cfg.bufferSize)
235
236 go func() {
237 defer close(out)
238 defer func() {
239 if cfg.observer != nil {
240 cfg.observer.OnDone()
241 }
242 }()
243
244 count := 0
245 for count < n {
246 select {
247 case v, ok := <-in:
248 if !ok {
249 return
250 }
251 if cfg.observer != nil {
252 cfg.observer.OnItem()
253 }
254 select {
255 case out <- v:
256 count++
257 case <-ctx.Done():
258 return
259 }
260 case <-ctx.Done():
261 return
262 }
263 }
264 }()
265
266 return out
267 }
268}
269
270// Sink consumes items with bounded concurrency without forwarding them.
271func Sink[T any](concurrency int, fn func(context.Context, T), opts ...StageOption) Stage[T] {
272 if concurrency < 1 {
273 panic("pipeline: Sink concurrency must be at least 1")
274 }
275 cfg := applyOpts(opts)
276
277 return func(ctx context.Context, in <-chan T) <-chan T {
278 done := make(chan T)
279 var wg sync.WaitGroup
280
281 for range concurrency {
282 wg.Go(func() {
283 for {
284 select {
285 case v, ok := <-in:
286 if !ok {
287 return
288 }
289 fn(ctx, v)
290 if cfg.observer != nil {
291 cfg.observer.OnItem()
292 }
293 case <-ctx.Done():
294 return
295 }
296 }
297 })
298 }
299
300 go func() {
301 wg.Wait()
302 close(done)
303 if cfg.observer != nil {
304 cfg.observer.OnDone()
305 }
306 }()
307
308 return done
309 }
310}
diff --git a/pipeline/stage_test.go b/pipeline/stage_test.go
new file mode 100644
index 0000000..4579158
--- /dev/null
+++ b/pipeline/stage_test.go
@@ -0,0 +1,308 @@
1package pipeline
2
3import (
4 "context"
5 "sort"
6 "sync"
7 "sync/atomic"
8 "testing"
9)
10
11func TestMap_Basic(t *testing.T) {
12 t.Parallel()
13
14 input := make(chan int, 3)
15 input <- 1
16 input <- 2
17 input <- 3
18 close(input)
19
20 out := Map(1, func(_ context.Context, v int) (int, bool) {
21 return v * 10, true
22 })(t.Context(), input)
23
24 var got []int
25 for v := range out {
26 got = append(got, v)
27 }
28
29 sort.Ints(got)
30 if len(got) != 3 || got[0] != 10 || got[1] != 20 || got[2] != 30 {
31 t.Fatalf("expected [10 20 30], got %v", got)
32 }
33}
34
35func TestMap_Concurrent(t *testing.T) {
36 t.Parallel()
37
38 input := make(chan int, 100)
39 for i := range 100 {
40 input <- i
41 }
42 close(input)
43
44 out := Map(5, func(_ context.Context, v int) (int, bool) {
45 return v, true
46 })(t.Context(), input)
47
48 var got []int
49 for v := range out {
50 got = append(got, v)
51 }
52
53 if len(got) != 100 {
54 t.Fatalf("expected 100 items, got %d", len(got))
55 }
56
57 sort.Ints(got)
58 for i, v := range got {
59 if v != i {
60 t.Fatalf("missing item %d, got %v at position %d", i, v, i)
61 }
62 }
63}
64
65func TestMap_Drop(t *testing.T) {
66 t.Parallel()
67
68 input := make(chan int, 5)
69 for i := 1; i <= 5; i++ {
70 input <- i
71 }
72 close(input)
73
74 out := Map(1, func(_ context.Context, v int) (int, bool) {
75 return v, v > 3
76 })(t.Context(), input)
77
78 var got []int
79 for v := range out {
80 got = append(got, v)
81 }
82
83 if len(got) != 2 {
84 t.Fatalf("expected 2 items, got %d: %v", len(got), got)
85 }
86}
87
88func TestMap_PanicOnZeroConcurrency(t *testing.T) {
89 t.Parallel()
90
91 defer func() {
92 if r := recover(); r == nil {
93 t.Error("expected panic for concurrency < 1")
94 }
95 }()
96
97 Map(0, func(_ context.Context, v int) (int, bool) { return v, true })
98}
99
100func TestFilter(t *testing.T) {
101 t.Parallel()
102
103 input := make(chan int, 10)
104 for i := 1; i <= 10; i++ {
105 input <- i
106 }
107 close(input)
108
109 out := Filter(func(v int) bool { return v%3 == 0 })(t.Context(), input)
110
111 var got []int
112 for v := range out {
113 got = append(got, v)
114 }
115
116 if len(got) != 3 {
117 t.Fatalf("expected [3 6 9], got %v", got)
118 }
119}
120
121func TestForEach(t *testing.T) {
122 t.Parallel()
123
124 input := make(chan int, 3)
125 input <- 1
126 input <- 2
127 input <- 3
128 close(input)
129
130 var sum atomic.Int32
131 out := ForEach(func(v int) {
132 sum.Add(int32(v))
133 })(t.Context(), input)
134
135 var got []int
136 for v := range out {
137 got = append(got, v)
138 }
139
140 if sum.Load() != 6 {
141 t.Fatalf("expected side effect sum 6, got %d", sum.Load())
142 }
143 if len(got) != 3 {
144 t.Fatalf("expected 3 forwarded items, got %d", len(got))
145 }
146}
147
148func TestBuffer(t *testing.T) {
149 t.Parallel()
150
151 input := make(chan int, 5)
152 for i := 1; i <= 5; i++ {
153 input <- i
154 }
155 close(input)
156
157 out := Buffer[int](10)(t.Context(), input)
158
159 var got []int
160 for v := range out {
161 got = append(got, v)
162 }
163
164 if len(got) != 5 {
165 t.Fatalf("expected 5 items, got %d", len(got))
166 }
167}
168
169func TestBuffer_PanicOnZero(t *testing.T) {
170 t.Parallel()
171
172 defer func() {
173 if r := recover(); r == nil {
174 t.Error("expected panic for buffer size < 1")
175 }
176 }()
177
178 Buffer[int](0)
179}
180
181func TestTee(t *testing.T) {
182 t.Parallel()
183
184 input := make(chan int, 3)
185 input <- 1
186 input <- 2
187 input <- 3
188 close(input)
189
190 side := make(chan int, 10)
191 out := Tee(side)(t.Context(), input)
192
193 var got []int
194 for v := range out {
195 got = append(got, v)
196 }
197 close(side)
198
199 var sideGot []int
200 for v := range side {
201 sideGot = append(sideGot, v)
202 }
203
204 if len(got) != 3 {
205 t.Fatalf("expected 3 forwarded items, got %d", len(got))
206 }
207 if len(sideGot) != 3 {
208 t.Fatalf("expected 3 side items, got %d", len(sideGot))
209 }
210}
211
212func TestTee_NilSide(t *testing.T) {
213 t.Parallel()
214
215 defer func() {
216 if r := recover(); r == nil {
217 t.Error("expected panic for nil side channel")
218 }
219 }()
220
221 Tee[int](nil)
222}
223
224func TestTake(t *testing.T) {
225 t.Parallel()
226 ctx, cancel := context.WithCancel(t.Context())
227 defer cancel()
228
229 input := make(chan int, 10)
230 for i := 1; i <= 10; i++ {
231 input <- i
232 }
233 close(input)
234
235 out := Take[int](3)(ctx, input)
236
237 var got []int
238 for v := range out {
239 got = append(got, v)
240 }
241
242 if len(got) != 3 {
243 t.Fatalf("expected 3 items, got %d: %v", len(got), got)
244 }
245}
246
247func TestTake_Zero(t *testing.T) {
248 t.Parallel()
249
250 input := make(chan int, 5)
251 for i := 1; i <= 5; i++ {
252 input <- i
253 }
254 close(input)
255
256 out := Take[int](0)(t.Context(), input)
257
258 count := 0
259 for range out {
260 count++
261 }
262
263 if count != 0 {
264 t.Fatalf("expected 0 items from Take(0), got %d", count)
265 }
266}
267
268func TestSink(t *testing.T) {
269 t.Parallel()
270
271 input := make(chan int, 5)
272 for i := 1; i <= 5; i++ {
273 input <- i
274 }
275 close(input)
276
277 var mu sync.Mutex
278 var collected []int
279
280 out := Sink(2, func(_ context.Context, v int) {
281 mu.Lock()
282 collected = append(collected, v)
283 mu.Unlock()
284 })(t.Context(), input)
285
286 for range out {
287 t.Fatal("Sink should not forward items")
288 }
289
290 mu.Lock()
291 defer mu.Unlock()
292
293 if len(collected) != 5 {
294 t.Fatalf("expected 5 items consumed, got %d", len(collected))
295 }
296}
297
298func TestSink_PanicOnZeroConcurrency(t *testing.T) {
299 t.Parallel()
300
301 defer func() {
302 if r := recover(); r == nil {
303 t.Error("expected panic for concurrency < 1")
304 }
305 }()
306
307 Sink(0, func(_ context.Context, v int) {})
308}