github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/scheduler/scheduler_test.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/prometheus/client_golang/prometheus"
     8  	"github.com/prometheus/client_golang/prometheus/promauto"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  	"google.golang.org/grpc/metadata"
    12  
    13  	"github.com/grafana/loki/pkg/scheduler/schedulerpb"
    14  	util_log "github.com/grafana/loki/pkg/util/log"
    15  )
    16  
    17  func TestScheduler_setRunState(t *testing.T) {
    18  
    19  	// This test is a bit crude, the method is not the most directly testable but
    20  	// this covers us to make sure we don't accidentally change the behavior of
    21  	// the little bit of logic which runs/stops the scheduler and makes sure we
    22  	// send a shutdown message to disconnect frontends.
    23  
    24  	// To avoid a lot more complicated test setup of calling NewScheduler instead
    25  	// we make a Scheduler with the things required to avoid nil pointers
    26  	s := Scheduler{
    27  		log: util_log.Logger,
    28  		schedulerRunning: promauto.With(prometheus.DefaultRegisterer).NewGauge(prometheus.GaugeOpts{
    29  			Name: "cortex_query_scheduler_running",
    30  			Help: "Value will be 1 if the scheduler is in the ReplicationSet and actively receiving/processing requests",
    31  		}),
    32  	}
    33  	mock := &mockSchedulerForFrontendFrontendLoopServer{}
    34  	s.connectedFrontends = map[string]*connectedFrontend{
    35  		"127.0.0.1:9095": {
    36  			connections: 0,
    37  			frontend:    mock,
    38  			ctx:         nil,
    39  			cancel:      nil,
    40  		},
    41  	}
    42  
    43  	// not_running, shouldRun == false
    44  	assert.False(t, s.shouldRun.Load())
    45  
    46  	// not_running -> running, shouldRun == true
    47  	s.setRunState(true)
    48  	assert.True(t, s.shouldRun.Load())
    49  
    50  	// running -> running, shouldRun == true
    51  	s.setRunState(true)
    52  	assert.True(t, s.shouldRun.Load())
    53  
    54  	// running -> not_running, shouldRun == false, shutdown message sent
    55  	s.setRunState(false)
    56  	assert.False(t, s.shouldRun.Load())
    57  	assert.Equal(t, schedulerpb.SHUTTING_DOWN, mock.msg.Status)
    58  	mock.msg = nil
    59  
    60  	// not_running -> not_running, shouldRun == false, no shutdown message sent
    61  	s.setRunState(false)
    62  	assert.Nil(t, mock.msg)
    63  
    64  }
    65  
    66  type mockSchedulerForFrontendFrontendLoopServer struct {
    67  	msg *schedulerpb.SchedulerToFrontend
    68  }
    69  
    70  func (m *mockSchedulerForFrontendFrontendLoopServer) Send(frontend *schedulerpb.SchedulerToFrontend) error {
    71  	m.msg = frontend
    72  	return nil
    73  }
    74  
    75  func (m mockSchedulerForFrontendFrontendLoopServer) Recv() (*schedulerpb.FrontendToScheduler, error) {
    76  	panic("implement me")
    77  }
    78  
    79  func (m mockSchedulerForFrontendFrontendLoopServer) SetHeader(md metadata.MD) error {
    80  	panic("implement me")
    81  }
    82  
    83  func (m mockSchedulerForFrontendFrontendLoopServer) SendHeader(md metadata.MD) error {
    84  	panic("implement me")
    85  }
    86  
    87  func (m mockSchedulerForFrontendFrontendLoopServer) SetTrailer(md metadata.MD) {
    88  	panic("implement me")
    89  }
    90  
    91  func (m mockSchedulerForFrontendFrontendLoopServer) Context() context.Context {
    92  	panic("implement me")
    93  }
    94  
    95  func (m mockSchedulerForFrontendFrontendLoopServer) SendMsg(msg interface{}) error {
    96  	panic("implement me")
    97  }
    98  
    99  func (m mockSchedulerForFrontendFrontendLoopServer) RecvMsg(msg interface{}) error {
   100  	panic("implement me")
   101  }