github.com/prebid/prebid-server/v2@v2.18.0/util/task/ticker_task_test.go (about)

     1  package task_test
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/prebid/prebid-server/v2/util/task"
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  type MockRunner struct {
    13  	ExpectationMet chan struct{}
    14  	actualCalls    int
    15  	expectedCalls  int
    16  	mutex          sync.Mutex
    17  }
    18  
    19  func NewMockRunner(expectedCalls int) *MockRunner {
    20  	return &MockRunner{
    21  		ExpectationMet: make(chan struct{}),
    22  		expectedCalls:  expectedCalls,
    23  	}
    24  }
    25  
    26  func (m *MockRunner) Run() error {
    27  	m.mutex.Lock()
    28  	defer m.mutex.Unlock()
    29  	m.actualCalls++
    30  
    31  	if m.expectedCalls == m.actualCalls {
    32  		close(m.ExpectationMet)
    33  	}
    34  
    35  	return nil
    36  }
    37  
    38  func (m *MockRunner) RunCount() int {
    39  	m.mutex.Lock()
    40  	defer m.mutex.Unlock()
    41  	return m.actualCalls
    42  }
    43  
    44  func TestStartWithSingleRun(t *testing.T) {
    45  	// Setup Initial Run Only:
    46  	expectedRuns := 1
    47  	runner := NewMockRunner(expectedRuns)
    48  	interval := 0 * time.Millisecond // forces a single run
    49  	ticker := task.NewTickerTask(interval, runner)
    50  
    51  	// Execute:
    52  	ticker.Start()
    53  
    54  	// Verify:
    55  	select {
    56  	case <-runner.ExpectationMet:
    57  	case <-time.After(250 * time.Millisecond):
    58  		assert.Failf(t, "Runner Calls", "expected %v calls, observed %v calls", expectedRuns, runner.RunCount())
    59  	}
    60  
    61  	// Verify No Additional Runs:
    62  	time.Sleep(50 * time.Millisecond)
    63  	assert.Equal(t, expectedRuns, runner.RunCount(), "runner should not run after Stop is called")
    64  }
    65  
    66  func TestStartWithPeriodicRun(t *testing.T) {
    67  	// Setup Initial Run + One Periodic Run:
    68  	expectedRuns := 2
    69  	runner := NewMockRunner(expectedRuns)
    70  	interval := 10 * time.Millisecond
    71  	ticker := task.NewTickerTask(interval, runner)
    72  
    73  	// Execute:
    74  	ticker.Start()
    75  
    76  	// Verify Expected Runs:
    77  	select {
    78  	case <-runner.ExpectationMet:
    79  		ticker.Stop()
    80  	case <-time.After(250 * time.Millisecond):
    81  		assert.Failf(t, "Runner Calls", "expected %v calls, observed %v calls", expectedRuns, runner.RunCount())
    82  	}
    83  
    84  	// Verify No Additional Runs After Stop:
    85  	time.Sleep(50 * time.Millisecond)
    86  	assert.Equal(t, expectedRuns, runner.RunCount(), "runner should not run after Stop is called")
    87  }