113 lines
3.9 KiB
Go
113 lines
3.9 KiB
Go
package drivers
|
|
|
|
import (
|
|
"io/ioutil"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/mitchellh/go-testing-interface"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/net/context"
|
|
|
|
plugin "github.com/hashicorp/go-plugin"
|
|
"github.com/hashicorp/nomad/helper/testlog"
|
|
"github.com/hashicorp/nomad/plugins/base"
|
|
"github.com/hashicorp/nomad/plugins/shared/hclspec"
|
|
)
|
|
|
|
type DriverHarness struct {
|
|
DriverPlugin
|
|
client *plugin.GRPCClient
|
|
server *plugin.GRPCServer
|
|
t testing.T
|
|
}
|
|
|
|
func NewDriverHarness(t testing.T, d DriverPlugin) *DriverHarness {
|
|
client, server := plugin.TestPluginGRPCConn(t, map[string]plugin.Plugin{
|
|
base.PluginTypeDriver: &PluginDriver{
|
|
impl: d,
|
|
logger: testlog.HCLogger(t),
|
|
},
|
|
})
|
|
|
|
raw, err := client.Dispense(base.PluginTypeDriver)
|
|
if err != nil {
|
|
t.Fatalf("err dispensing plugin: %v", err)
|
|
}
|
|
|
|
dClient := raw.(DriverPlugin)
|
|
h := &DriverHarness{
|
|
client: client,
|
|
server: server,
|
|
DriverPlugin: dClient,
|
|
}
|
|
|
|
return h
|
|
}
|
|
|
|
func (h *DriverHarness) Kill() {
|
|
h.client.Close()
|
|
h.server.Stop()
|
|
}
|
|
|
|
// MkAllocDir creates a tempory directory and allocdir structure.
|
|
// A cleanup func is returned and should be defered so as to not leak dirs
|
|
// between tests.
|
|
func (h *DriverHarness) MkAllocDir(t *TaskConfig) func() {
|
|
allocDir, err := ioutil.TempDir("", "nomad_driver_harness-")
|
|
require.NoError(h.t, err)
|
|
require.NoError(h.t, os.Mkdir(filepath.Join(allocDir, t.Name), os.ModePerm))
|
|
require.NoError(h.t, os.MkdirAll(filepath.Join(allocDir, "alloc/logs"), os.ModePerm))
|
|
t.AllocDir = allocDir
|
|
return func() { os.RemoveAll(allocDir) }
|
|
}
|
|
|
|
// MockDriver is used for testing.
|
|
// Each function can be set as a closure to make assertions about how data
|
|
// is passed through the base plugin layer.
|
|
type MockDriver struct {
|
|
base.MockPlugin
|
|
TaskConfigSchemaF func() (*hclspec.Spec, error)
|
|
FingerprintF func(context.Context) (<-chan *Fingerprint, error)
|
|
CapabilitiesF func() (*Capabilities, error)
|
|
RecoverTaskF func(*TaskHandle) error
|
|
StartTaskF func(*TaskConfig) (*TaskHandle, error)
|
|
WaitTaskF func(context.Context, string) (<-chan *ExitResult, error)
|
|
StopTaskF func(string, time.Duration, string) error
|
|
DestroyTaskF func(string, bool) error
|
|
InspectTaskF func(string) (*TaskStatus, error)
|
|
TaskStatsF func(string) (*TaskStats, error)
|
|
TaskEventsF func(context.Context) (<-chan *TaskEvent, error)
|
|
SignalTaskF func(string, string) error
|
|
ExecTaskF func(string, []string, time.Duration) (*ExecTaskResult, error)
|
|
}
|
|
|
|
func (d *MockDriver) TaskConfigSchema() (*hclspec.Spec, error) { return d.TaskConfigSchemaF() }
|
|
func (d *MockDriver) Fingerprint(ctx context.Context) (<-chan *Fingerprint, error) {
|
|
return d.FingerprintF(ctx)
|
|
}
|
|
func (d *MockDriver) Capabilities() (*Capabilities, error) { return d.CapabilitiesF() }
|
|
func (d *MockDriver) RecoverTask(h *TaskHandle) error { return d.RecoverTaskF(h) }
|
|
func (d *MockDriver) StartTask(c *TaskConfig) (*TaskHandle, error) { return d.StartTaskF(c) }
|
|
func (d *MockDriver) WaitTask(ctx context.Context, id string) (<-chan *ExitResult, error) {
|
|
return d.WaitTaskF(ctx, id)
|
|
}
|
|
func (d *MockDriver) StopTask(taskID string, timeout time.Duration, signal string) error {
|
|
return d.StopTaskF(taskID, timeout, signal)
|
|
}
|
|
func (d *MockDriver) DestroyTask(taskID string, force bool) error {
|
|
return d.DestroyTaskF(taskID, force)
|
|
}
|
|
func (d *MockDriver) InspectTask(taskID string) (*TaskStatus, error) { return d.InspectTaskF(taskID) }
|
|
func (d *MockDriver) TaskStats(taskID string) (*TaskStats, error) { return d.TaskStats(taskID) }
|
|
func (d *MockDriver) TaskEvents(ctx context.Context) (<-chan *TaskEvent, error) {
|
|
return d.TaskEventsF(ctx)
|
|
}
|
|
func (d *MockDriver) SignalTask(taskID string, signal string) error {
|
|
return d.SignalTask(taskID, signal)
|
|
}
|
|
func (d *MockDriver) ExecTask(taskID string, cmd []string, timeout time.Duration) (*ExecTaskResult, error) {
|
|
return d.ExecTaskF(taskID, cmd, timeout)
|
|
}
|