github.com/jrxfive/nomad@v0.6.1-0.20170802162750-1fef470e89bf/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/consul/api"
    12  	"github.com/hashicorp/nomad/client/driver"
    13  	cstructs "github.com/hashicorp/nomad/client/structs"
    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 != "checks" {
    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  	// checksFn allows injecting return values for the Checks function.
    48  	checksFn func(*structs.Allocation) ([]*api.AgentCheck, error)
    49  }
    50  
    51  func newMockConsulServiceClient() *mockConsulServiceClient {
    52  	m := mockConsulServiceClient{
    53  		ops:    make([]mockConsulOp, 0, 20),
    54  		logger: log.New(ioutil.Discard, "", 0),
    55  	}
    56  	if testing.Verbose() {
    57  		m.logger = log.New(os.Stderr, "", log.LstdFlags)
    58  	}
    59  	return &m
    60  }
    61  
    62  func (m *mockConsulServiceClient) UpdateTask(allocID string, old, new *structs.Task, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) error {
    63  	m.mu.Lock()
    64  	defer m.mu.Unlock()
    65  	m.logger.Printf("[TEST] mock_consul: UpdateTask(%q, %v, %v, %T, %x)", allocID, old, new, exec, net.Hash())
    66  	m.ops = append(m.ops, newMockConsulOp("update", allocID, new, exec, net))
    67  	return nil
    68  }
    69  
    70  func (m *mockConsulServiceClient) RegisterTask(allocID string, task *structs.Task, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) error {
    71  	m.mu.Lock()
    72  	defer m.mu.Unlock()
    73  	m.logger.Printf("[TEST] mock_consul: RegisterTask(%q, %q, %T, %x)", allocID, task.Name, exec, net.Hash())
    74  	m.ops = append(m.ops, newMockConsulOp("add", allocID, task, exec, net))
    75  	return nil
    76  }
    77  
    78  func (m *mockConsulServiceClient) RemoveTask(allocID string, task *structs.Task) {
    79  	m.mu.Lock()
    80  	defer m.mu.Unlock()
    81  	m.logger.Printf("[TEST] mock_consul: RemoveTask(%q, %q)", allocID, task.Name)
    82  	m.ops = append(m.ops, newMockConsulOp("remove", allocID, task, nil, nil))
    83  }
    84  
    85  func (m *mockConsulServiceClient) Checks(alloc *structs.Allocation) ([]*api.AgentCheck, error) {
    86  	m.mu.Lock()
    87  	defer m.mu.Unlock()
    88  	m.logger.Printf("[TEST] mock_consul: Checks(%q)", alloc.ID)
    89  	m.ops = append(m.ops, newMockConsulOp("checks", alloc.ID, nil, nil, nil))
    90  
    91  	if m.checksFn != nil {
    92  		return m.checksFn(alloc)
    93  	}
    94  
    95  	return nil, nil
    96  }