open-nomad/client/consul/consul_testing.go
Seth Hoenig 93cf770edb client: enable nomad client to request and set SI tokens for tasks
When a job is configured with Consul Connect aware tasks (i.e. sidecar),
the Nomad Client should be able to request from Consul (through Nomad Server)
Service Identity tokens specific to those tasks.
2020-01-31 19:03:38 -06:00

114 lines
3.2 KiB
Go

package consul
import (
"fmt"
"sync"
"time"
log "github.com/hashicorp/go-hclog"
"github.com/hashicorp/nomad/command/agent/consul"
testing "github.com/mitchellh/go-testing-interface"
)
// MockConsulOp represents the register/deregister operations.
type MockConsulOp struct {
Op string // add, remove, or update
AllocID string
Name string // task or group name
OccurredAt time.Time
}
func NewMockConsulOp(op, allocID, name string) MockConsulOp {
switch op {
case "add", "remove", "update", "alloc_registrations",
"add_group", "remove_group", "update_group", "update_ttl":
default:
panic(fmt.Errorf("invalid consul op: %s", op))
}
return MockConsulOp{
Op: op,
AllocID: allocID,
Name: name,
OccurredAt: time.Now(),
}
}
// MockConsulServiceClient implements the ConsulServiceAPI interface to record
// and log task registration/deregistration.
type MockConsulServiceClient struct {
ops []MockConsulOp
mu sync.Mutex
logger log.Logger
// AllocRegistrationsFn allows injecting return values for the
// AllocRegistrations function.
AllocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error)
}
func NewMockConsulServiceClient(t testing.T, logger log.Logger) *MockConsulServiceClient {
logger = logger.Named("mock_consul")
m := MockConsulServiceClient{
ops: make([]MockConsulOp, 0, 20),
logger: logger,
}
return &m
}
func (m *MockConsulServiceClient) UpdateWorkload(old, newSvcs *consul.WorkloadServices) error {
m.mu.Lock()
defer m.mu.Unlock()
m.logger.Trace("UpdateWorkload", "alloc_id", newSvcs.AllocID, "name", newSvcs.Name(),
"old_services", len(old.Services), "new_services", len(newSvcs.Services),
)
m.ops = append(m.ops, NewMockConsulOp("update", newSvcs.AllocID, newSvcs.Name()))
return nil
}
func (m *MockConsulServiceClient) RegisterWorkload(svcs *consul.WorkloadServices) error {
m.mu.Lock()
defer m.mu.Unlock()
m.logger.Trace("RegisterWorkload", "alloc_id", svcs.AllocID, "name", svcs.Name(),
"services", len(svcs.Services),
)
m.ops = append(m.ops, NewMockConsulOp("add", svcs.AllocID, svcs.Name()))
return nil
}
func (m *MockConsulServiceClient) RemoveWorkload(svcs *consul.WorkloadServices) {
m.mu.Lock()
defer m.mu.Unlock()
m.logger.Trace("RemoveWorkload", "alloc_id", svcs.AllocID, "name", svcs.Name(),
"services", len(svcs.Services),
)
m.ops = append(m.ops, NewMockConsulOp("remove", svcs.AllocID, svcs.Name()))
}
func (m *MockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) {
m.mu.Lock()
defer m.mu.Unlock()
m.logger.Trace("AllocRegistrations", "alloc_id", allocID)
m.ops = append(m.ops, NewMockConsulOp("alloc_registrations", allocID, ""))
if m.AllocRegistrationsFn != nil {
return m.AllocRegistrationsFn(allocID)
}
return nil, nil
}
func (m *MockConsulServiceClient) UpdateTTL(checkID, output, status string) error {
// TODO(tgross): this method is here so we can implement the
// interface but the locking we need for testing creates a lot
// of opportunities for deadlocks in testing that will never
// appear in live code.
m.logger.Trace("UpdateTTL", "check_id", checkID, "status", status)
return nil
}
func (m *MockConsulServiceClient) GetOps() []MockConsulOp {
m.mu.Lock()
defer m.mu.Unlock()
return m.ops
}