github.com/keybase/client/go@v0.0.0-20240309051027-028f7c731f8b/kbfs/libcontext/delayed_cancellation_test.go (about)

     1  // Copyright 2016 Keybase Inc. All rights reserved.
     2  // Use of this source code is governed by a BSD
     3  // license that can be found in the LICENSE file.
     4  
     5  package libcontext
     6  
     7  import (
     8  	"testing"
     9  	"time"
    10  
    11  	"golang.org/x/net/context"
    12  )
    13  
    14  type testDCKeyType int
    15  
    16  const (
    17  	testDCKey testDCKeyType = iota
    18  )
    19  
    20  func TestReplayableContext(t *testing.T) {
    21  	t.Parallel()
    22  
    23  	ctx := context.Background()
    24  	ctx = NewContextReplayable(ctx, func(ctx context.Context) context.Context {
    25  		return context.WithValue(ctx, testDCKey, "O_O")
    26  	})
    27  	ctx, cancel := context.WithCancel(ctx)
    28  
    29  	ctx, err := NewContextWithReplayFrom(ctx)
    30  	if err != nil {
    31  		t.Fatalf("calling NewContextWithReplayFrom error: %s", err)
    32  	}
    33  
    34  	// Test if replay was run properly
    35  	if ctx.Value(testDCKey) != "O_O" {
    36  		t.Fatalf("NewContextWithReplayFrom did not replay attached replayFunc")
    37  	}
    38  
    39  	// Test if cancellation is disabled
    40  	cancel()
    41  	select {
    42  	case <-ctx.Done():
    43  		t.Fatalf("NewContextWithReplayFrom did not disconnect the cancel function")
    44  	default:
    45  	}
    46  
    47  	// make sure the new ctx is also replayable
    48  	ctx, err = NewContextWithReplayFrom(ctx)
    49  	if err != nil {
    50  		t.Fatalf("calling NewContextWithReplayFrom error: %s", err)
    51  	}
    52  	if ctx.Value(testDCKey) != "O_O" {
    53  		t.Fatalf("NewContextWithReplayFrom did not replay attached replayFunc")
    54  	}
    55  }
    56  
    57  func makeContextWithDelayedCancellation(t *testing.T) (
    58  	ctx context.Context, originalCancel context.CancelFunc) {
    59  	ctx = context.Background()
    60  	ctx = NewContextReplayable(ctx, func(ctx context.Context) context.Context {
    61  		return context.WithValue(ctx, testDCKey, "O_O")
    62  	})
    63  	ctx, cancel := context.WithCancel(ctx)
    64  
    65  	ctx, err := NewContextWithCancellationDelayer(ctx)
    66  	if err != nil {
    67  		t.Fatalf("calling NewContextWithCancellationDelayer error: %s", err)
    68  	}
    69  
    70  	// Test NewContextWithCancellationDelayer does replay properly
    71  	if ctx.Value(testDCKey) != "O_O" {
    72  		t.Fatalf(
    73  			"NewContextWithCancellationDelayer did not replay attached replayFunc")
    74  	}
    75  
    76  	return ctx, cancel
    77  }
    78  
    79  func TestDelayedCancellationCancelWhileNotEnabled(t *testing.T) {
    80  	t.Parallel()
    81  
    82  	ctx, cancel := makeContextWithDelayedCancellation(t)
    83  
    84  	cancel()
    85  
    86  	select {
    87  	case <-ctx.Done():
    88  	case <-time.After(100 * time.Millisecond):
    89  		t.Fatalf("Cancellation did not happen even though " +
    90  			"EnableDelayedCancellationWithGracePeriod has not been called yet")
    91  	}
    92  }
    93  
    94  func TestDelayedCancellationCleanupWhileNotEnabled(t *testing.T) {
    95  	t.Parallel()
    96  
    97  	ctx, _ := makeContextWithDelayedCancellation(t)
    98  
    99  	if err := CleanupCancellationDelayer(ctx); err != nil {
   100  		t.Fatalf("calling CleanupCancellationDelayer error: %s", err)
   101  	}
   102  
   103  	select {
   104  	case <-ctx.Done():
   105  	case <-time.After(100 * time.Millisecond):
   106  		t.Fatalf("Cancellation did not happen even though " +
   107  			"EnableDelayedCancellationWithGracePeriod has not been called yet")
   108  	}
   109  }
   110  
   111  func TestDelayedCancellationSecondEnable(t *testing.T) {
   112  	t.Parallel()
   113  
   114  	ctx, cancel := makeContextWithDelayedCancellation(t)
   115  	defer cancel()
   116  
   117  	err := EnableDelayedCancellationWithGracePeriod(ctx, 0)
   118  	if err != nil {
   119  		t.Fatalf("1st EnableDelayedCancellationWithGracePeriod failed: %v", err)
   120  	}
   121  	cancel()
   122  	<-ctx.Done()
   123  	// parent context is not canceled; second "enable" should succeed even it's
   124  	// after grace period
   125  	err = EnableDelayedCancellationWithGracePeriod(ctx, 0)
   126  	if err == nil {
   127  		t.Fatalf("2nd EnableDelayedCancellationWithGracePeriod succeeded even " +
   128  			"though more than grace period has passed since parent context was " +
   129  			"canceled")
   130  	}
   131  }
   132  
   133  func TestDelayedCancellationEnabled(t *testing.T) {
   134  	t.Parallel()
   135  
   136  	ctx, cancel := makeContextWithDelayedCancellation(t)
   137  	err := EnableDelayedCancellationWithGracePeriod(ctx, 50*time.Millisecond)
   138  	if err != nil {
   139  		t.Fatalf("EnableDelayedCancellationWithGracePeriod failed: %v", err)
   140  	}
   141  
   142  	cancel()
   143  
   144  	select {
   145  	case <-ctx.Done():
   146  		t.Fatalf("Cancellation is not delayed")
   147  	case <-time.After(10 * time.Millisecond):
   148  	}
   149  
   150  	<-ctx.Done()
   151  
   152  	// if test timeouts, then it's a failure: Cancellation did not happen after
   153  	// grace period
   154  }