pull_test.go

     1  // Copyright 2023 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  //go:build goexperiment.rangefunc
     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