github.com/bigcommerce/nomad@v0.9.3-bc/client/allocrunner/testing.go (about)

     1  // +build !release
     2  
     3  package allocrunner
     4  
     5  import (
     6  	"sync"
     7  	"testing"
     8  
     9  	"github.com/hashicorp/nomad/client/allocwatcher"
    10  	clientconfig "github.com/hashicorp/nomad/client/config"
    11  	"github.com/hashicorp/nomad/client/consul"
    12  	"github.com/hashicorp/nomad/client/devicemanager"
    13  	"github.com/hashicorp/nomad/client/pluginmanager/drivermanager"
    14  	"github.com/hashicorp/nomad/client/state"
    15  	"github.com/hashicorp/nomad/client/vaultclient"
    16  	"github.com/hashicorp/nomad/nomad/structs"
    17  	"github.com/stretchr/testify/require"
    18  )
    19  
    20  // MockStateUpdater implements the AllocStateHandler interface and records
    21  // alloc updates.
    22  type MockStateUpdater struct {
    23  	Updates []*structs.Allocation
    24  	mu      sync.Mutex
    25  }
    26  
    27  // AllocStateUpdated implements the AllocStateHandler interface and records an
    28  // alloc update.
    29  func (m *MockStateUpdater) AllocStateUpdated(alloc *structs.Allocation) {
    30  	m.mu.Lock()
    31  	m.Updates = append(m.Updates, alloc)
    32  	m.mu.Unlock()
    33  }
    34  
    35  // Last returns a copy of the last alloc (or nil) update. Safe for concurrent
    36  // access with updates.
    37  func (m *MockStateUpdater) Last() *structs.Allocation {
    38  	m.mu.Lock()
    39  	defer m.mu.Unlock()
    40  	n := len(m.Updates)
    41  	if n == 0 {
    42  		return nil
    43  	}
    44  	return m.Updates[n-1].Copy()
    45  }
    46  
    47  // Reset resets the recorded alloc updates.
    48  func (m *MockStateUpdater) Reset() {
    49  	m.mu.Lock()
    50  	m.Updates = nil
    51  	m.mu.Unlock()
    52  }
    53  
    54  func testAllocRunnerConfig(t *testing.T, alloc *structs.Allocation) (*Config, func()) {
    55  	clientConf, cleanup := clientconfig.TestClientConfig(t)
    56  	conf := &Config{
    57  		// Copy the alloc in case the caller edits and reuses it
    58  		Alloc:              alloc.Copy(),
    59  		Logger:             clientConf.Logger,
    60  		ClientConfig:       clientConf,
    61  		StateDB:            state.NoopDB{},
    62  		Consul:             consul.NewMockConsulServiceClient(t, clientConf.Logger),
    63  		Vault:              vaultclient.NewMockVaultClient(),
    64  		StateUpdater:       &MockStateUpdater{},
    65  		PrevAllocWatcher:   allocwatcher.NoopPrevAlloc{},
    66  		PrevAllocMigrator:  allocwatcher.NoopPrevAlloc{},
    67  		DeviceManager:      devicemanager.NoopMockManager(),
    68  		DriverManager:      drivermanager.TestDriverManager(t),
    69  		ServersContactedCh: make(chan struct{}),
    70  	}
    71  	return conf, cleanup
    72  }
    73  
    74  func TestAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation) (*allocRunner, func()) {
    75  	t.Helper()
    76  	cfg, cleanup := testAllocRunnerConfig(t, alloc)
    77  	ar, err := NewAllocRunner(cfg)
    78  	if err != nil {
    79  		require.NoError(t, err, "Failed to setup AllocRunner")
    80  	}
    81  
    82  	return ar, cleanup
    83  }