github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/worker/frontend_processor_test.go (about)

     1  package worker
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/go-kit/log"
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  	"go.uber.org/atomic"
    12  	"google.golang.org/grpc"
    13  
    14  	"github.com/cortexproject/cortex/pkg/util/test"
    15  )
    16  
    17  func TestRecvFailDoesntCancelProcess(t *testing.T) {
    18  	ctx, cancel := context.WithCancel(context.Background())
    19  	defer cancel()
    20  
    21  	// We use random port here, hopefully without any gRPC server.
    22  	cc, err := grpc.DialContext(ctx, "localhost:999", grpc.WithInsecure())
    23  	require.NoError(t, err)
    24  
    25  	cfg := Config{}
    26  	mgr := newFrontendProcessor(cfg, nil, log.NewNopLogger())
    27  	running := atomic.NewBool(false)
    28  	go func() {
    29  		running.Store(true)
    30  		defer running.Store(false)
    31  
    32  		mgr.processQueriesOnSingleStream(ctx, cc, "test:12345")
    33  	}()
    34  
    35  	test.Poll(t, time.Second, true, func() interface{} {
    36  		return running.Load()
    37  	})
    38  
    39  	// Wait a bit, and verify that processQueriesOnSingleStream is still running, and hasn't stopped
    40  	// just because it cannot contact frontend.
    41  	time.Sleep(100 * time.Millisecond)
    42  	assert.Equal(t, true, running.Load())
    43  
    44  	cancel()
    45  	test.Poll(t, time.Second, false, func() interface{} {
    46  		return running.Load()
    47  	})
    48  }
    49  
    50  func TestContextCancelStopsProcess(t *testing.T) {
    51  	ctx, cancel := context.WithCancel(context.Background())
    52  	defer cancel()
    53  
    54  	// We use random port here, hopefully without any gRPC server.
    55  	cc, err := grpc.DialContext(ctx, "localhost:999", grpc.WithInsecure())
    56  	require.NoError(t, err)
    57  
    58  	pm := newProcessorManager(ctx, &mockProcessor{}, cc, "test")
    59  	pm.concurrency(1)
    60  
    61  	test.Poll(t, time.Second, 1, func() interface{} {
    62  		return int(pm.currentProcessors.Load())
    63  	})
    64  
    65  	cancel()
    66  
    67  	test.Poll(t, time.Second, 0, func() interface{} {
    68  		return int(pm.currentProcessors.Load())
    69  	})
    70  
    71  	pm.stop()
    72  	test.Poll(t, time.Second, 0, func() interface{} {
    73  		return int(pm.currentProcessors.Load())
    74  	})
    75  }