1 files changed, 118 insertions, 0 deletions
diff --git a/src/iter/pull_test.go b/src/iter/pull_test.go
new file mode 100644
index 0000000000..38e0ee993a
--- /dev/null
+++ b/src/iter/pull_test.go
@@ -0,0 +1,118 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build goexperiment.rangefunc
+
+package iter
+
+import (
+	"fmt"
+	"runtime"
+	"testing"
+)
+
+func count(n int) Seq[int] {
+	return func(yield func(int) bool) {
+		for i := range n {
+			if !yield(i) {
+				break
+			}
+		}
+	}
+}
+
+func squares(n int) Seq2[int, int64] {
+	return func(yield func(int, int64) bool) {
+		for i := range n {
+			if !yield(i, int64(i)*int64(i)) {
+				break
+			}
+		}
+	}
+}
+
+func TestPull(t *testing.T) {
+
+	for end := 0; end <= 3; end++ {
+		t.Run(fmt.Sprint(end), func(t *testing.T) {
+			ng := runtime.NumGoroutine()
+			wantNG := func(want int) {
+				if xg := runtime.NumGoroutine() - ng; xg != want {
+					t.Helper()
+					t.Errorf("have %d extra goroutines, want %d", xg, want)
+				}
+			}
+			wantNG(0)
+			next, stop := Pull(count(3))
+			wantNG(1)
+			for i := range end {
+				v, ok := next()
+				if v != i || ok != true {
+					t.Fatalf("next() = %d, %v, want %d, %v", v, ok, i, true)
+				}
+				wantNG(1)
+			}
+			wantNG(1)
+			if end < 3 {
+				stop()
+				wantNG(0)
+			}
+			for range 2 {
+				v, ok := next()
+				if v != 0 || ok != false {
+					t.Fatalf("next() = %d, %v, want %d, %v", v, ok, 0, false)
+				}
+				wantNG(0)
+			}
+			wantNG(0)
+
+			stop()
+			stop()
+			stop()
+			wantNG(0)
+		})
+	}
+}
+
+func TestPull2(t *testing.T) {
+	for end := 0; end <= 3; end++ {
+		t.Run(fmt.Sprint(end), func(t *testing.T) {
+			ng := runtime.NumGoroutine()
+			wantNG := func(want int) {
+				if xg := runtime.NumGoroutine() - ng; xg != want {
+					t.Helper()
+					t.Errorf("have %d extra goroutines, want %d", xg, want)
+				}
+			}
+			wantNG(0)
+			next, stop := Pull2(squares(3))
+			wantNG(1)
+			for i := range end {
+				k, v, ok := next()
+				if k != i || v != int64(i*i) || ok != true {
+					t.Fatalf("next() = %d, %d, %v, want %d, %d, %v", k, v, ok, i, i*i, true)
+				}
+				wantNG(1)
+			}
+			wantNG(1)
+			if end < 3 {
+				stop()
+				wantNG(0)
+			}
+			for range 2 {
+				k, v, ok := next()
+				if v != 0 || ok != false {
+					t.Fatalf("next() = %d, %d, %v, want %d, %d, %v", k, v, ok, 0, 0, false)
+				}
+				wantNG(0)
+			}
+			wantNG(0)
+
+			stop()
+			stop()
+			stop()
+			wantNG(0)
+		})
+	}
+}