github.com/djenriquez/nomad-1@v0.8.1/client/consul_testing.go (about)

     1  package client
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"sync"
     7  
     8  	"github.com/hashicorp/nomad/client/driver"
     9  	cstructs "github.com/hashicorp/nomad/client/structs"
    10  	"github.com/hashicorp/nomad/command/agent/consul"
    11  	"github.com/hashicorp/nomad/helper/testlog"
    12  	"github.com/hashicorp/nomad/nomad/structs"
    13  	"github.com/mitchellh/go-testing-interface"
    14  )
    15  
    16  // mockConsulOp represents the register/deregister operations.
    17  type mockConsulOp struct {
    18  	op      string // add, remove, or update
    19  	allocID string
    20  	task    *structs.Task
    21  	exec    driver.ScriptExecutor
    22  	net     *cstructs.DriverNetwork
    23  }
    24  
    25  func newMockConsulOp(op, allocID string, task *structs.Task, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) mockConsulOp {
    26  	if op != "add" && op != "remove" && op != "update" && op != "alloc_registrations" {
    27  		panic(fmt.Errorf("invalid consul op: %s", op))
    28  	}
    29  	return mockConsulOp{
    30  		op:      op,
    31  		allocID: allocID,
    32  		task:    task,
    33  		exec:    exec,
    34  		net:     net,
    35  	}
    36  }
    37  
    38  // mockConsulServiceClient implements the ConsulServiceAPI interface to record
    39  // and log task registration/deregistration.
    40  type mockConsulServiceClient struct {
    41  	ops []mockConsulOp
    42  	mu  sync.Mutex
    43  
    44  	logger *log.Logger
    45  
    46  	// allocRegistrationsFn allows injecting return values for the
    47  	// AllocRegistrations function.
    48  	allocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error)
    49  }
    50  
    51  func newMockConsulServiceClient(t testing.T) *mockConsulServiceClient {
    52  	m := mockConsulServiceClient{
    53  		ops:    make([]mockConsulOp, 0, 20),
    54  		logger: testlog.Logger(t),
    55  	}
    56  	return &m
    57  }
    58  
    59  func (m *mockConsulServiceClient) UpdateTask(allocID string, old, new *structs.Task, restarter consul.TaskRestarter, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) error {
    60  	m.mu.Lock()
    61  	defer m.mu.Unlock()
    62  	m.logger.Printf("[TEST] mock_consul: UpdateTask(%q, %v, %v, %T, %x)", allocID, old, new, exec, net.Hash())
    63  	m.ops = append(m.ops, newMockConsulOp("update", allocID, new, exec, net))
    64  	return nil
    65  }
    66  
    67  func (m *mockConsulServiceClient) RegisterTask(allocID string, task *structs.Task, restarter consul.TaskRestarter, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) error {
    68  	m.mu.Lock()
    69  	defer m.mu.Unlock()
    70  	m.logger.Printf("[TEST] mock_consul: RegisterTask(%q, %q, %T, %x)", allocID, task.Name, exec, net.Hash())
    71  	m.ops = append(m.ops, newMockConsulOp("add", allocID, task, exec, net))
    72  	return nil
    73  }
    74  
    75  func (m *mockConsulServiceClient) RemoveTask(allocID string, task *structs.Task) {
    76  	m.mu.Lock()
    77  	defer m.mu.Unlock()
    78  	m.logger.Printf("[TEST] mock_consul: RemoveTask(%q, %q)", allocID, task.Name)
    79  	m.ops = append(m.ops, newMockConsulOp("remove", allocID, task, nil, nil))
    80  }
    81  
    82  func (m *mockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) {
    83  	m.mu.Lock()
    84  	defer m.mu.Unlock()
    85  	m.logger.Printf("[TEST] mock_consul: AllocRegistrations(%q)", allocID)
    86  	m.ops = append(m.ops, newMockConsulOp("alloc_registrations", allocID, nil, nil, nil))
    87  
    88  	if m.allocRegistrationsFn != nil {
    89  		return m.allocRegistrationsFn(allocID)
    90  	}
    91  
    92  	return nil, nil
    93  }