github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/common/fscontext/fscontext_test.go (about)

     1  // Copyright 2024 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package fscontext_test
    16  
    17  import (
    18  	"context"
    19  	"errors"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/google/fleetspeak/fleetspeak/src/common/fscontext"
    24  )
    25  
    26  var (
    27  	errMagic = errors.New("magic")
    28  	errDone  = errors.New("done")
    29  )
    30  
    31  var shortDuration = 100 * time.Millisecond
    32  
    33  func TestWithDoneChanNotCanceled(t *testing.T) {
    34  	// Given
    35  	doneCh := make(chan struct{})
    36  	defer close(doneCh)
    37  
    38  	ctx, cancel := fscontext.WithDoneChan(context.TODO(), errDone, doneCh)
    39  	defer cancel(nil)
    40  
    41  	// When nothing happens...
    42  
    43  	// Then
    44  	select {
    45  	case <-ctx.Done():
    46  		t.Errorf("Expected ctx not canceled, got: %v", context.Cause(ctx))
    47  	case <-time.After(shortDuration):
    48  		t.Log("not canceled after", shortDuration, "- looks good")
    49  	}
    50  }
    51  
    52  func TestWithDoneChanCanceledThroughChannel(t *testing.T) {
    53  	// Given
    54  	doneCh := make(chan struct{})
    55  
    56  	ctx, cancel := fscontext.WithDoneChan(context.TODO(), errDone, doneCh)
    57  	defer cancel(nil)
    58  
    59  	// When
    60  	close(doneCh)
    61  
    62  	// Then
    63  	select {
    64  	case <-time.After(shortDuration):
    65  		t.Errorf("timeout waiting for context cancelation")
    66  	case <-ctx.Done():
    67  		if err := ctx.Err(); err != context.Canceled {
    68  			t.Errorf("done channel closed: ctx.Err() = %v, want canceled", err)
    69  		}
    70  		if !errors.Is(context.Cause(ctx), errDone) {
    71  			t.Errorf("done channel closed: context.Cause(ctx) = %v, want errors.Is(…, %v)", context.Cause(ctx), errDone)
    72  		}
    73  	}
    74  }
    75  
    76  func TestWithDoneChanCanceledThroughOuterContext(t *testing.T) {
    77  	// Given
    78  	doneCh := make(chan struct{})
    79  	defer close(doneCh)
    80  
    81  	ctx, cancel1 := context.WithCancelCause(context.TODO())
    82  	defer cancel1(nil)
    83  	ctx, cancel2 := fscontext.WithDoneChan(ctx, errDone, doneCh)
    84  	defer cancel2(nil)
    85  
    86  	// When
    87  	cancel1(errMagic)
    88  
    89  	// Then
    90  	select {
    91  	case <-time.After(shortDuration):
    92  		t.Errorf("timeout waiting for context cancelation")
    93  	case <-ctx.Done():
    94  		if err := ctx.Err(); err != context.Canceled {
    95  			t.Errorf("done channel closed: ctx.Err() = %v, want canceled", err)
    96  		}
    97  		if !errors.Is(context.Cause(ctx), errMagic) {
    98  			t.Errorf("done channel closed: context.Cause(ctx) = %v, want errors.Is(…, %v)", context.Cause(ctx), errMagic)
    99  		}
   100  	}
   101  }
   102  
   103  func TestWithDoneChanCanceledThroughOwnContext(t *testing.T) {
   104  	// Given
   105  	doneCh := make(chan struct{})
   106  	defer close(doneCh)
   107  
   108  	ctx, cancel := fscontext.WithDoneChan(context.TODO(), errDone, doneCh)
   109  	defer cancel(nil)
   110  
   111  	// When
   112  	cancel(errMagic)
   113  
   114  	// Then
   115  	select {
   116  	case <-time.After(shortDuration):
   117  		t.Errorf("timeout waiting for context cancelation")
   118  	case <-ctx.Done():
   119  		if err := ctx.Err(); err != context.Canceled {
   120  			t.Errorf("done channel closed: ctx.Err() = %v, want canceled", err)
   121  		}
   122  		if !errors.Is(context.Cause(ctx), errMagic) {
   123  			t.Errorf("done channel closed: context.Cause(ctx) = %v, want errors.Is(…, %v)", context.Cause(ctx), errMagic)
   124  		}
   125  	}
   126  }