github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/allocrunner/testing.go (about)

     1  //go:build !release
     2  // +build !release
     3  
     4  package allocrunner
     5  
     6  import (
     7  	"fmt"
     8  	"sync"
     9  	"testing"
    10  
    11  	"github.com/hashicorp/nomad/client/allocrunner/taskrunner/getter"
    12  	"github.com/hashicorp/nomad/client/allocwatcher"
    13  	clientconfig "github.com/hashicorp/nomad/client/config"
    14  	"github.com/hashicorp/nomad/client/consul"
    15  	"github.com/hashicorp/nomad/client/devicemanager"
    16  	"github.com/hashicorp/nomad/client/lib/cgutil"
    17  	"github.com/hashicorp/nomad/client/pluginmanager/drivermanager"
    18  	"github.com/hashicorp/nomad/client/serviceregistration/checks/checkstore"
    19  	"github.com/hashicorp/nomad/client/serviceregistration/mock"
    20  	"github.com/hashicorp/nomad/client/serviceregistration/wrapper"
    21  	"github.com/hashicorp/nomad/client/state"
    22  	"github.com/hashicorp/nomad/client/vaultclient"
    23  	"github.com/hashicorp/nomad/nomad/structs"
    24  	"github.com/hashicorp/nomad/testutil"
    25  	"github.com/stretchr/testify/require"
    26  )
    27  
    28  // MockStateUpdater implements the AllocStateHandler interface and records
    29  // alloc updates.
    30  type MockStateUpdater struct {
    31  	Updates []*structs.Allocation
    32  	mu      sync.Mutex
    33  }
    34  
    35  // AllocStateUpdated implements the AllocStateHandler interface and records an
    36  // alloc update.
    37  func (m *MockStateUpdater) AllocStateUpdated(alloc *structs.Allocation) {
    38  	m.mu.Lock()
    39  	m.Updates = append(m.Updates, alloc)
    40  	m.mu.Unlock()
    41  }
    42  
    43  // PutAllocation satisfies the AllocStateHandler interface.
    44  func (m *MockStateUpdater) PutAllocation(alloc *structs.Allocation) (err error) {
    45  	return
    46  }
    47  
    48  // Last returns a copy of the last alloc (or nil) update. Safe for concurrent
    49  // access with updates.
    50  func (m *MockStateUpdater) Last() *structs.Allocation {
    51  	m.mu.Lock()
    52  	defer m.mu.Unlock()
    53  	n := len(m.Updates)
    54  	if n == 0 {
    55  		return nil
    56  	}
    57  	return m.Updates[n-1].Copy()
    58  }
    59  
    60  // Reset resets the recorded alloc updates.
    61  func (m *MockStateUpdater) Reset() {
    62  	m.mu.Lock()
    63  	m.Updates = nil
    64  	m.mu.Unlock()
    65  }
    66  
    67  func testAllocRunnerConfig(t *testing.T, alloc *structs.Allocation) (*Config, func()) {
    68  	clientConf, cleanup := clientconfig.TestClientConfig(t)
    69  
    70  	consulRegMock := mock.NewServiceRegistrationHandler(clientConf.Logger)
    71  	nomadRegMock := mock.NewServiceRegistrationHandler(clientConf.Logger)
    72  
    73  	stateDB := new(state.NoopDB)
    74  
    75  	conf := &Config{
    76  		// Copy the alloc in case the caller edits and reuses it
    77  		Alloc:              alloc.Copy(),
    78  		Logger:             clientConf.Logger,
    79  		ClientConfig:       clientConf,
    80  		StateDB:            stateDB,
    81  		Consul:             consulRegMock,
    82  		ConsulSI:           consul.NewMockServiceIdentitiesClient(),
    83  		Vault:              vaultclient.NewMockVaultClient(),
    84  		StateUpdater:       &MockStateUpdater{},
    85  		PrevAllocWatcher:   allocwatcher.NoopPrevAlloc{},
    86  		PrevAllocMigrator:  allocwatcher.NoopPrevAlloc{},
    87  		DeviceManager:      devicemanager.NoopMockManager(),
    88  		DriverManager:      drivermanager.TestDriverManager(t),
    89  		CpusetManager:      new(cgutil.NoopCpusetManager),
    90  		ServersContactedCh: make(chan struct{}),
    91  		ServiceRegWrapper:  wrapper.NewHandlerWrapper(clientConf.Logger, consulRegMock, nomadRegMock),
    92  		CheckStore:         checkstore.NewStore(clientConf.Logger, stateDB),
    93  		Getter:             getter.TestSandbox(t),
    94  	}
    95  
    96  	return conf, cleanup
    97  }
    98  
    99  func TestAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation) (*allocRunner, func()) {
   100  	t.Helper()
   101  	cfg, cleanup := testAllocRunnerConfig(t, alloc)
   102  	ar, err := NewAllocRunner(cfg)
   103  	if err != nil {
   104  		require.NoError(t, err, "Failed to setup AllocRunner")
   105  	}
   106  
   107  	return ar, cleanup
   108  }
   109  
   110  func WaitForClientState(t *testing.T, ar *allocRunner, state string) {
   111  	testutil.WaitForResult(func() (bool, error) {
   112  		got := ar.AllocState().ClientStatus
   113  		return got == state,
   114  			fmt.Errorf("expected alloc runner to be in state %s, got %s", state, got)
   115  	}, func(err error) {
   116  		require.NoError(t, err)
   117  	})
   118  }