github.com/banmanh482/nomad@v0.11.8/helper/pluginutils/loader/plugin_test.go (about) 1 package loader 2 3 import ( 4 "context" 5 "flag" 6 "fmt" 7 "os" 8 "testing" 9 "time" 10 11 log "github.com/hashicorp/go-hclog" 12 plugin "github.com/hashicorp/go-plugin" 13 "github.com/hashicorp/nomad/plugins/base" 14 "github.com/hashicorp/nomad/plugins/device" 15 "github.com/hashicorp/nomad/plugins/shared/hclspec" 16 ) 17 18 type stringSliceFlags []string 19 20 func (i *stringSliceFlags) String() string { 21 return "my string representation" 22 } 23 24 func (i *stringSliceFlags) Set(value string) error { 25 *i = append(*i, value) 26 return nil 27 } 28 29 // TestMain runs either the tests or runs a mock plugin based on the passed 30 // flags 31 func TestMain(m *testing.M) { 32 var plugin, configSchema bool 33 var name, pluginType, pluginVersion string 34 var pluginApiVersions stringSliceFlags 35 flag.BoolVar(&plugin, "plugin", false, "run binary as a plugin") 36 flag.BoolVar(&configSchema, "config-schema", true, "return a config schema") 37 flag.StringVar(&name, "name", "", "plugin name") 38 flag.StringVar(&pluginType, "type", "", "plugin type") 39 flag.StringVar(&pluginVersion, "version", "", "plugin version") 40 flag.Var(&pluginApiVersions, "api-version", "supported plugin API version") 41 flag.Parse() 42 43 if plugin { 44 if err := pluginMain(name, pluginType, pluginVersion, pluginApiVersions, configSchema); err != nil { 45 fmt.Println(err.Error()) 46 os.Exit(1) 47 } 48 } else { 49 os.Exit(m.Run()) 50 } 51 } 52 53 // pluginMain starts a mock plugin using the passed parameters 54 func pluginMain(name, pluginType, version string, apiVersions []string, config bool) error { 55 // Validate passed parameters 56 if name == "" || pluginType == "" { 57 return fmt.Errorf("name and plugin type must be specified") 58 } 59 60 switch pluginType { 61 case base.PluginTypeDevice: 62 default: 63 return fmt.Errorf("unsupported plugin type %q", pluginType) 64 } 65 66 // Create the mock plugin 67 m := &mockPlugin{ 68 name: name, 69 ptype: pluginType, 70 version: version, 71 apiVersions: apiVersions, 72 configSchema: config, 73 } 74 75 // Build the plugin map 76 pmap := map[string]plugin.Plugin{ 77 base.PluginTypeBase: &base.PluginBase{Impl: m}, 78 } 79 switch pluginType { 80 case base.PluginTypeDevice: 81 pmap[base.PluginTypeDevice] = &device.PluginDevice{Impl: m} 82 } 83 84 // Serve the plugin 85 plugin.Serve(&plugin.ServeConfig{ 86 HandshakeConfig: base.Handshake, 87 Plugins: pmap, 88 GRPCServer: plugin.DefaultGRPCServer, 89 }) 90 91 return nil 92 } 93 94 // mockFactory returns a PluginFactory method which creates the mock plugin with 95 // the passed parameters 96 func mockFactory(name, ptype, version string, apiVersions []string, configSchema bool) func(context.Context, log.Logger) interface{} { 97 return func(ctx context.Context, log log.Logger) interface{} { 98 return &mockPlugin{ 99 name: name, 100 ptype: ptype, 101 version: version, 102 apiVersions: apiVersions, 103 configSchema: configSchema, 104 } 105 } 106 } 107 108 // mockPlugin is a plugin that meets various plugin interfaces but is only 109 // useful for testing. 110 type mockPlugin struct { 111 name string 112 ptype string 113 version string 114 apiVersions []string 115 configSchema bool 116 117 // config is built on SetConfig 118 config *mockPluginConfig 119 120 // nomadconfig is set on SetConfig 121 nomadConfig *base.AgentConfig 122 123 // negotiatedApiVersion is the version of the api to use and is set on 124 // SetConfig 125 negotiatedApiVersion string 126 } 127 128 // mockPluginConfig is the configuration for the mock plugin 129 type mockPluginConfig struct { 130 Foo string `codec:"foo"` 131 Bar int `codec:"bar"` 132 133 // ResKey is a key that is populated in the Env map when a device is 134 // reserved. 135 ResKey string `codec:"res_key"` 136 } 137 138 // PluginInfo returns the plugin information based on the passed fields when 139 // building the mock plugin 140 func (m *mockPlugin) PluginInfo() (*base.PluginInfoResponse, error) { 141 return &base.PluginInfoResponse{ 142 Type: m.ptype, 143 PluginApiVersions: m.apiVersions, 144 PluginVersion: m.version, 145 Name: m.name, 146 }, nil 147 } 148 149 func (m *mockPlugin) ConfigSchema() (*hclspec.Spec, error) { 150 if !m.configSchema { 151 return nil, nil 152 } 153 154 // configSpec is the hclspec for parsing the mock's configuration 155 configSpec := hclspec.NewObject(map[string]*hclspec.Spec{ 156 "foo": hclspec.NewAttr("foo", "string", false), 157 "bar": hclspec.NewAttr("bar", "number", false), 158 "res_key": hclspec.NewAttr("res_key", "string", false), 159 }) 160 161 return configSpec, nil 162 } 163 164 // SetConfig decodes the configuration and stores it 165 func (m *mockPlugin) SetConfig(c *base.Config) error { 166 var config mockPluginConfig 167 if len(c.PluginConfig) != 0 { 168 if err := base.MsgPackDecode(c.PluginConfig, &config); err != nil { 169 return err 170 } 171 } 172 173 m.config = &config 174 m.nomadConfig = c.AgentConfig 175 m.negotiatedApiVersion = c.ApiVersion 176 return nil 177 } 178 179 func (m *mockPlugin) Fingerprint(ctx context.Context) (<-chan *device.FingerprintResponse, error) { 180 return make(chan *device.FingerprintResponse), nil 181 } 182 183 func (m *mockPlugin) Reserve(deviceIDs []string) (*device.ContainerReservation, error) { 184 if m.config == nil || m.config.ResKey == "" { 185 return nil, nil 186 } 187 188 return &device.ContainerReservation{ 189 Envs: map[string]string{m.config.ResKey: "config-set"}, 190 }, nil 191 } 192 193 func (m *mockPlugin) Stats(ctx context.Context, interval time.Duration) (<-chan *device.StatsResponse, error) { 194 return make(chan *device.StatsResponse), nil 195 }