func (h *transportHandle) ServiceGenerator() ServiceGenerator { if !h.Running.Load() { panic(fmt.Sprintf("handle for plugin %q has already been closed", h.name)) } if _, hasFeature := h.Features[api.FeatureServiceGenerator]; !hasFeature { return nil } return &serviceGenerator{ handle: h, Running: h.Running, ServiceGenerator: api.NewServiceGeneratorClient(multiplex.NewClient( "ServiceGenerator", envelope.NewClient(_proto, h.Transport), )), } }
func TestServiceGenerator(t *testing.T) { transport, done := fakeEnvelopeClient() defer done() mockCtrl := gomock.NewController(t) defer mockCtrl.Finish() serviceGenerator := plugintest.NewMockServiceGenerator(mockCtrl) go Main(&Plugin{ Name: "hello", ServiceGenerator: serviceGenerator, }) pluginClient := api.NewPluginClient(multiplex.NewClient("Plugin", transport)) defer pluginClient.Goodbye() handshake, err := pluginClient.Handshake(&api.HandshakeRequest{}) require.NoError(t, err) assert.Equal(t, api.APIVersion, handshake.APIVersion) assert.Equal(t, version.Version, *handshake.LibraryVersion) assert.Equal(t, "hello", handshake.Name) assert.Contains(t, handshake.Features, api.FeatureServiceGenerator) sgClient := api.NewServiceGeneratorClient(multiplex.NewClient("ServiceGenerator", transport)) req := &api.GenerateServiceRequest{ RootServices: []api.ServiceID{1}, Services: map[api.ServiceID]*api.Service{ 1: { Name: "MyService", ThriftName: "MyService", Functions: []*api.Function{}, ParentID: (*api.ServiceID)(ptr.Int32(2)), ModuleID: 1, }, 2: { Name: "BaseService", ThriftName: "BaseService", Functions: []*api.Function{ { Name: "Healthy", ThriftName: "healthy", Arguments: []*api.Argument{}, }, }, ModuleID: 1, }, }, Modules: map[api.ModuleID]*api.Module{ 1: { ImportPath: "go.uber.org/thriftrw/plugin/fake", Directory: "fake", }, }, } res := &api.GenerateServiceResponse{ Files: map[string][]byte{ "fake/myservice/foo.go": {1, 2, 3}, "fake/baseservice/bar.go": {4, 5, 6}, "fake/baz.go": {7, 8, 9}, }, } serviceGenerator.EXPECT().Generate(req).Return(res, nil) gotRes, err := sgClient.Generate(req) if assert.NoError(t, err) { assert.Equal(t, res, gotRes) } }