github.com/zhizhiboom/nomad@v0.8.5-0.20180907175415-f28fd3a1a056/client/allocrunner/testing.go (about)

     1  package allocrunner
     2  
     3  import (
     4  	"io/ioutil"
     5  	"os"
     6  	"sync"
     7  	"testing"
     8  
     9  	"github.com/boltdb/bolt"
    10  	"github.com/hashicorp/nomad/client/config"
    11  	consulApi "github.com/hashicorp/nomad/client/consul"
    12  	"github.com/hashicorp/nomad/client/vaultclient"
    13  	"github.com/hashicorp/nomad/helper/testlog"
    14  	"github.com/hashicorp/nomad/nomad/mock"
    15  	"github.com/hashicorp/nomad/nomad/structs"
    16  )
    17  
    18  type MockAllocStateUpdater struct {
    19  	Allocs []*structs.Allocation
    20  	mu     sync.Mutex
    21  }
    22  
    23  // Update fulfills the TaskStateUpdater interface
    24  func (m *MockAllocStateUpdater) Update(alloc *structs.Allocation) {
    25  	m.mu.Lock()
    26  	m.Allocs = append(m.Allocs, alloc)
    27  	m.mu.Unlock()
    28  }
    29  
    30  // Last returns a copy of the last alloc (or nil) sync'd
    31  func (m *MockAllocStateUpdater) Last() *structs.Allocation {
    32  	m.mu.Lock()
    33  	defer m.mu.Unlock()
    34  	n := len(m.Allocs)
    35  	if n == 0 {
    36  		return nil
    37  	}
    38  	return m.Allocs[n-1].Copy()
    39  }
    40  
    41  func TestAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation, restarts bool) (*MockAllocStateUpdater, *AllocRunner) {
    42  	conf := config.DefaultConfig()
    43  	conf.Node = mock.Node()
    44  	conf.StateDir = os.TempDir()
    45  	conf.AllocDir = os.TempDir()
    46  	tmp, _ := ioutil.TempFile("", "state-db")
    47  	db, _ := bolt.Open(tmp.Name(), 0600, nil)
    48  	upd := &MockAllocStateUpdater{}
    49  	if !restarts {
    50  		*alloc.Job.LookupTaskGroup(alloc.TaskGroup).RestartPolicy = structs.RestartPolicy{Attempts: 0}
    51  		alloc.Job.Type = structs.JobTypeBatch
    52  	}
    53  	vclient := vaultclient.NewMockVaultClient()
    54  	ar := NewAllocRunner(testlog.Logger(t), conf, db, upd.Update, alloc, vclient, consulApi.NewMockConsulServiceClient(t), NoopPrevAlloc{})
    55  	return upd, ar
    56  }
    57  
    58  func TestAllocRunner(t *testing.T, restarts bool) (*MockAllocStateUpdater, *AllocRunner) {
    59  	// Use mock driver
    60  	alloc := mock.Alloc()
    61  	task := alloc.Job.TaskGroups[0].Tasks[0]
    62  	task.Driver = "mock_driver"
    63  	task.Config["run_for"] = "500ms"
    64  	return TestAllocRunnerFromAlloc(t, alloc, restarts)
    65  }