mirror of
https://github.com/kemko/nomad.git
synced 2026-01-03 17:05:43 +03:00
143 lines
2.8 KiB
Go
143 lines
2.8 KiB
Go
// Copyright (c) HashiCorp, Inc.
|
|
// SPDX-License-Identifier: BUSL-1.1
|
|
|
|
package testing
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/hashicorp/nomad/nomad/structs"
|
|
)
|
|
|
|
// MockEmitter is a mock of the EventEmitter interface.
|
|
type MockEmitter struct {
|
|
lock sync.Mutex
|
|
events []*structs.TaskEvent
|
|
}
|
|
|
|
func (m *MockEmitter) EmitEvent(ev *structs.TaskEvent) {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
m.events = append(m.events, ev)
|
|
}
|
|
|
|
func (m *MockEmitter) Events() []*structs.TaskEvent {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
return m.events
|
|
}
|
|
|
|
// MockTaskHooks is a mock of the TaskHooks interface useful for testing
|
|
type MockTaskHooks struct {
|
|
lock sync.Mutex
|
|
|
|
RestartCh chan struct{}
|
|
restarts int
|
|
|
|
SignalCh chan struct{}
|
|
signals []string
|
|
|
|
// SignalError is returned when Signal is called on the mock hook
|
|
SignalError error
|
|
|
|
UnblockCh chan struct{}
|
|
|
|
KillCh chan *structs.TaskEvent
|
|
killEvent *structs.TaskEvent
|
|
|
|
EmitEventCh chan *structs.TaskEvent
|
|
events []*structs.TaskEvent
|
|
|
|
// HasHandle can be set to simulate restoring a task after client restart
|
|
HasHandle bool
|
|
}
|
|
|
|
func NewMockTaskHooks() *MockTaskHooks {
|
|
return &MockTaskHooks{
|
|
UnblockCh: make(chan struct{}, 1),
|
|
RestartCh: make(chan struct{}, 1),
|
|
SignalCh: make(chan struct{}, 1),
|
|
KillCh: make(chan *structs.TaskEvent, 1),
|
|
EmitEventCh: make(chan *structs.TaskEvent, 1),
|
|
}
|
|
}
|
|
func (m *MockTaskHooks) Restart(ctx context.Context, event *structs.TaskEvent, failure bool) error {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
|
|
m.restarts++
|
|
select {
|
|
case m.RestartCh <- struct{}{}:
|
|
default:
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *MockTaskHooks) Signal(event *structs.TaskEvent, s string) error {
|
|
m.lock.Lock()
|
|
m.signals = append(m.signals, s)
|
|
m.lock.Unlock()
|
|
|
|
select {
|
|
case m.SignalCh <- struct{}{}:
|
|
default:
|
|
}
|
|
|
|
return m.SignalError
|
|
}
|
|
|
|
func (m *MockTaskHooks) Signals() []string {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
return m.signals
|
|
}
|
|
|
|
func (m *MockTaskHooks) Kill(ctx context.Context, event *structs.TaskEvent) error {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
|
|
m.killEvent = event
|
|
select {
|
|
case m.KillCh <- event:
|
|
default:
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *MockTaskHooks) IsRunning() bool {
|
|
return m.HasHandle
|
|
}
|
|
|
|
func (m *MockTaskHooks) EmitEvent(event *structs.TaskEvent) {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
|
|
m.events = append(m.events, event)
|
|
select {
|
|
case m.EmitEventCh <- event:
|
|
case <-m.EmitEventCh:
|
|
m.EmitEventCh <- event
|
|
}
|
|
}
|
|
|
|
func (m *MockTaskHooks) SetState(state string, event *structs.TaskEvent) {}
|
|
|
|
func (m *MockTaskHooks) KillEvent() *structs.TaskEvent {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
return m.killEvent
|
|
}
|
|
|
|
func (m *MockTaskHooks) Events() []*structs.TaskEvent {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
return m.events
|
|
}
|
|
|
|
func (m *MockTaskHooks) Restarts() int {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
return m.restarts
|
|
}
|