champa-client/pollingpacketconn_test.go

   1 package main
   2
   3 import (
   4         "bytes"
   5         "context"
   6         "errors"
   7         "io"
   8         "io/ioutil"
   9         "testing"
  10         "time"
  11
  12         "www.bamsoftware.com/git/champa.git/turbotunnel"
  13 )
  14
  15 // TestCloseCancelsPoll tests that calling Close cancels the context passed to
  16 // the poll function.
  17 func TestCloseCancelsPoll(t *testing.T) {
  18         beginCh := make(chan struct{})
  19         resultCh := make(chan error)
  20         // The poll function returns immediately with a nil error when its
  21         // context is canceled. It returns after a delay with a non-nil error if
  22         // its context is not canceled.
  23         var poll PollFunc = func(ctx context.Context, _ []byte) (io.ReadCloser, error) {
  24                 defer close(resultCh)
  25                 beginCh <- struct{}{}
  26                 select {
  27                 case <-ctx.Done():
  28                         resultCh <- nil
  29                 case <-time.After(5 * time.Second):
  30                         resultCh <- errors.New("poll was not canceled")
  31                 }
  32                 return ioutil.NopCloser(bytes.NewReader(nil)), nil
  33         }
  34         pconn := NewPollingPacketConn(turbotunnel.DummyAddr{}, poll)
  35         // Wait until the poll function has been called.
  36         <-beginCh
  37         // Close the connection.
  38         err := pconn.Close()
  39         if err != nil {
  40                 t.Fatal(err)
  41         }
  42         // Observe what happened inside the poll function. Closing the
  43         // connection should have canceled the context.
  44         err = <-resultCh
  45         if err != nil {
  46                 t.Fatal(err)
  47         }
  48 }