github.com/grafana/pyroscope@v1.18.0/pkg/metastore/compaction/scheduler/scheduler_test.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/hashicorp/raft"
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/mock"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1/raft_log"
    12  	"github.com/grafana/pyroscope/pkg/iter"
    13  	"github.com/grafana/pyroscope/pkg/test/mocks/mockscheduler"
    14  )
    15  
    16  func TestScheduler_UpdateSchedule(t *testing.T) {
    17  	store := new(mockscheduler.MockJobStore)
    18  	store.On("StoreJobPlan", mock.Anything, &raft_log.CompactionJobPlan{Name: "1"}).Return(nil).Once()
    19  	store.On("StoreJobState", mock.Anything, &raft_log.CompactionJobState{Name: "1"}).Return(nil).Once()
    20  	store.On("StoreJobState", mock.Anything, &raft_log.CompactionJobState{Name: "2"}).Return(nil).Once()
    21  	store.On("DeleteJobPlan", mock.Anything, "3").Return(nil).Once()
    22  	store.On("DeleteJobPlan", mock.Anything, "4").Return(nil).Once()
    23  	store.On("DeleteJobState", mock.Anything, "3").Return(nil).Once()
    24  	store.On("DeleteJobState", mock.Anything, "4").Return(nil).Once()
    25  
    26  	scheduler := NewScheduler(Config{}, store, nil)
    27  	for _, job := range []*raft_log.CompactionJobState{
    28  		{Name: "1"},
    29  		{Name: "2"},
    30  		{Name: "3"},
    31  		{Name: "4"},
    32  	} {
    33  		scheduler.queue.put(job)
    34  	}
    35  
    36  	update := &raft_log.CompactionPlanUpdate{
    37  		NewJobs: []*raft_log.NewCompactionJob{{
    38  			State: &raft_log.CompactionJobState{Name: "1"},
    39  			Plan:  &raft_log.CompactionJobPlan{Name: "1"},
    40  		}},
    41  		UpdatedJobs: []*raft_log.UpdatedCompactionJob{{
    42  			State: &raft_log.CompactionJobState{Name: "2"},
    43  		}},
    44  		CompletedJobs: []*raft_log.CompletedCompactionJob{{
    45  			State: &raft_log.CompactionJobState{Name: "3"},
    46  		}},
    47  		EvictedJobs: []*raft_log.EvictedCompactionJob{{
    48  			State: &raft_log.CompactionJobState{Name: "4"},
    49  		}},
    50  	}
    51  
    52  	require.NoError(t, scheduler.UpdateSchedule(nil, update))
    53  	s := scheduler.NewSchedule(nil, &raft.Log{Index: 3})
    54  
    55  	store.On("GetJobPlan", mock.Anything, "1").Return(new(raft_log.CompactionJobPlan), nil).Once()
    56  	assigment, err := s.AssignJob()
    57  	require.NoError(t, err)
    58  	assert.NotNil(t, assigment)
    59  
    60  	store.On("GetJobPlan", mock.Anything, "2").Return(new(raft_log.CompactionJobPlan), nil).Once()
    61  	assigment, err = s.AssignJob()
    62  	require.NoError(t, err)
    63  	assert.NotNil(t, assigment)
    64  
    65  	assigment, err = s.AssignJob()
    66  	require.NoError(t, err)
    67  	assert.Nil(t, assigment)
    68  
    69  	assert.Equal(t, jobQueuePop(scheduler.queue), update.NewJobs[0].State)
    70  	assert.Equal(t, jobQueuePop(scheduler.queue), update.UpdatedJobs[0].State)
    71  	assert.Nil(t, jobQueuePop(scheduler.queue))
    72  	newStatsCollector(scheduler).collectStats(func(level int, stats queueStats) {
    73  		assert.Equal(t, 0, level)
    74  		assert.Equal(t, stats, queueStats{
    75  			addedTotal:     4,
    76  			completedTotal: 1,
    77  			evictedTotal:   1,
    78  		})
    79  	})
    80  
    81  	store.AssertExpectations(t)
    82  }
    83  
    84  func TestScheduler_Restore(t *testing.T) {
    85  	store := new(mockscheduler.MockJobStore)
    86  	scheduler := NewScheduler(Config{}, store, nil)
    87  
    88  	store.On("ListEntries", mock.Anything).Return(iter.NewSliceIterator([]*raft_log.CompactionJobState{
    89  		{Name: "1", Token: 1},
    90  		{Name: "2", Token: 1},
    91  	}))
    92  
    93  	require.NoError(t, scheduler.Restore(nil))
    94  	s := scheduler.NewSchedule(nil, &raft.Log{Index: 3})
    95  
    96  	store.On("GetJobPlan", mock.Anything, "1").Return(new(raft_log.CompactionJobPlan), nil).Once()
    97  	assigment, err := s.AssignJob()
    98  	require.NoError(t, err)
    99  	assert.NotNil(t, assigment)
   100  
   101  	store.On("GetJobPlan", mock.Anything, "2").Return(new(raft_log.CompactionJobPlan), nil).Once()
   102  	assigment, err = s.AssignJob()
   103  	require.NoError(t, err)
   104  	assert.NotNil(t, assigment)
   105  
   106  	assigment, err = s.AssignJob()
   107  	require.NoError(t, err)
   108  	assert.Nil(t, assigment)
   109  
   110  	store.AssertExpectations(t)
   111  }