github.com/neatlab/neatio@v1.7.3-0.20220425043230-d903e92fcc75/utilities/event/subscription_test.go (about)

     1  package event
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  )
     9  
    10  var errInts = errors.New("error in subscribeInts")
    11  
    12  func subscribeInts(max, fail int, c chan<- int) Subscription {
    13  	return NewSubscription(func(quit <-chan struct{}) error {
    14  		for i := 0; i < max; i++ {
    15  			if i >= fail {
    16  				return errInts
    17  			}
    18  			select {
    19  			case c <- i:
    20  			case <-quit:
    21  				return nil
    22  			}
    23  		}
    24  		return nil
    25  	})
    26  }
    27  
    28  func TestNewSubscriptionError(t *testing.T) {
    29  	t.Parallel()
    30  
    31  	channel := make(chan int)
    32  	sub := subscribeInts(10, 2, channel)
    33  loop:
    34  	for want := 0; want < 10; want++ {
    35  		select {
    36  		case got := <-channel:
    37  			if got != want {
    38  				t.Fatalf("wrong int %d, want %d", got, want)
    39  			}
    40  		case err := <-sub.Err():
    41  			if err != errInts {
    42  				t.Fatalf("wrong error: got %q, want %q", err, errInts)
    43  			}
    44  			if want != 2 {
    45  				t.Fatalf("got errInts at int %d, should be received at 2", want)
    46  			}
    47  			break loop
    48  		}
    49  	}
    50  	sub.Unsubscribe()
    51  
    52  	err, ok := <-sub.Err()
    53  	if err != nil {
    54  		t.Fatal("got non-nil error after Unsubscribe")
    55  	}
    56  	if ok {
    57  		t.Fatal("channel still open after Unsubscribe")
    58  	}
    59  }
    60  
    61  func TestResubscribe(t *testing.T) {
    62  	t.Parallel()
    63  
    64  	var i int
    65  	nfails := 6
    66  	sub := Resubscribe(100*time.Millisecond, func(ctx context.Context) (Subscription, error) {
    67  
    68  		i++
    69  		if i == 2 {
    70  
    71  			time.Sleep(200 * time.Millisecond)
    72  		}
    73  		if i < nfails {
    74  			return nil, errors.New("oops")
    75  		}
    76  		sub := NewSubscription(func(unsubscribed <-chan struct{}) error { return nil })
    77  		return sub, nil
    78  	})
    79  
    80  	<-sub.Err()
    81  	if i != nfails {
    82  		t.Fatalf("resubscribe function called %d times, want %d times", i, nfails)
    83  	}
    84  }
    85  
    86  func TestResubscribeAbort(t *testing.T) {
    87  	t.Parallel()
    88  
    89  	done := make(chan error)
    90  	sub := Resubscribe(0, func(ctx context.Context) (Subscription, error) {
    91  		select {
    92  		case <-ctx.Done():
    93  			done <- nil
    94  		case <-time.After(2 * time.Second):
    95  			done <- errors.New("context given to resubscribe function not canceled within 2s")
    96  		}
    97  		return nil, nil
    98  	})
    99  
   100  	sub.Unsubscribe()
   101  	if err := <-done; err != nil {
   102  		t.Fatal(err)
   103  	}
   104  }