github.com/blixtra/nomad@v0.7.2-0.20171221000451-da9a1d7bb050/client/consul_test.go (about)

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