github.com/nibnait/go-learn@v0.0.0-20220227013611-dfa47ea6d2da/src/test/chapter/ch4/28_select_close_test.go (about)

     1  package ch4
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"testing"
     7  )
     8  
     9  func dataProducer(ch chan int, wg *sync.WaitGroup) {
    10  	go func() {
    11  		defer func() {
    12  			close(ch)
    13  		}()
    14  
    15  		for i := 0; i < 10; i++ {
    16  			ch <- i
    17  		}
    18  
    19  		wg.Done()
    20  	}()
    21  }
    22  
    23  func dataReceiver(ch chan int, wg *sync.WaitGroup) {
    24  	go func() {
    25  		for {
    26  			if data, ok := <-ch; ok {
    27  				fmt.Println(data)
    28  			} else {
    29  				break
    30  			}
    31  		}
    32  		wg.Done()
    33  	}()
    34  }
    35  
    36  func TestCloseChannel(t *testing.T) {
    37  	var wg sync.WaitGroup
    38  	ch := make(chan int)
    39  
    40  	wg.Add(1)
    41  	dataProducer(ch, &wg)
    42  
    43  	wg.Add(1)
    44  	dataReceiver(ch, &wg)
    45  	wg.Add(1)
    46  	dataReceiver(ch, &wg)
    47  
    48  	wg.Wait()
    49  }