Skip to content

Commit

Permalink
feature: multiplexing: fix unit tests (#14007)
Browse files Browse the repository at this point in the history
* fix grpc_server tests and add coverage

* update run_config tests

* add happy path test case for grpc_server ID from context

* update test helpers
  • Loading branch information
fairclothjm committed Feb 14, 2022
1 parent 4782d45 commit 25620a1
Show file tree
Hide file tree
Showing 3 changed files with 151 additions and 90 deletions.
4 changes: 0 additions & 4 deletions sdk/database/dbplugin/v5/grpc_server.go
Original file line number Diff line number Diff line change
Expand Up @@ -74,10 +74,6 @@ func (g gRPCServer) getDatabaseInternal(ctx context.Context) (Database, error) {
return nil, err
}

if id == "" {
return nil, fmt.Errorf("no instance ID found for multiplexed plugin")
}

if db, ok := g.instances[id]; ok {
return db, nil
}
Expand Down
136 changes: 94 additions & 42 deletions sdk/database/dbplugin/v5/grpc_server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,9 @@ import (
"github.com/golang/protobuf/ptypes"
"github.com/golang/protobuf/ptypes/timestamp"
"github.com/hashicorp/vault/sdk/database/dbplugin/v5/proto"
"github.com/hashicorp/vault/sdk/helper/pluginutil"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
)

Expand Down Expand Up @@ -77,14 +79,9 @@ func TestGRPCServer_Initialize(t *testing.T) {

for name, test := range tests {
t.Run(name, func(t *testing.T) {
g := gRPCServer{
impl: test.db,
}

// Context doesn't need to timeout since this is just passed through
ctx := context.Background()
idCtx, g := testGrpcServer(t, test.db)
resp, err := g.Initialize(idCtx, test.req)

resp, err := g.Initialize(ctx, test.req)
if test.expectErr && err == nil {
t.Fatalf("err expected, got nil")
}
Expand Down Expand Up @@ -252,14 +249,9 @@ func TestGRPCServer_NewUser(t *testing.T) {

for name, test := range tests {
t.Run(name, func(t *testing.T) {
g := gRPCServer{
impl: test.db,
}

// Context doesn't need to timeout since this is just passed through
ctx := context.Background()
idCtx, g := testGrpcServer(t, test.db)
resp, err := g.NewUser(idCtx, test.req)

resp, err := g.NewUser(ctx, test.req)
if test.expectErr && err == nil {
t.Fatalf("err expected, got nil")
}
Expand Down Expand Up @@ -362,14 +354,9 @@ func TestGRPCServer_UpdateUser(t *testing.T) {

for name, test := range tests {
t.Run(name, func(t *testing.T) {
g := gRPCServer{
impl: test.db,
}

// Context doesn't need to timeout since this is just passed through
ctx := context.Background()
idCtx, g := testGrpcServer(t, test.db)
resp, err := g.UpdateUser(idCtx, test.req)

resp, err := g.UpdateUser(ctx, test.req)
if test.expectErr && err == nil {
t.Fatalf("err expected, got nil")
}
Expand Down Expand Up @@ -430,14 +417,9 @@ func TestGRPCServer_DeleteUser(t *testing.T) {

for name, test := range tests {
t.Run(name, func(t *testing.T) {
g := gRPCServer{
impl: test.db,
}
idCtx, g := testGrpcServer(t, test.db)
resp, err := g.DeleteUser(idCtx, test.req)

// Context doesn't need to timeout since this is just passed through
ctx := context.Background()

resp, err := g.DeleteUser(ctx, test.req)
if test.expectErr && err == nil {
t.Fatalf("err expected, got nil")
}
Expand Down Expand Up @@ -488,14 +470,9 @@ func TestGRPCServer_Type(t *testing.T) {

for name, test := range tests {
t.Run(name, func(t *testing.T) {
g := gRPCServer{
impl: test.db,
}

// Context doesn't need to timeout since this is just passed through
ctx := context.Background()
idCtx, g := testGrpcServer(t, test.db)
resp, err := g.Type(idCtx, &proto.Empty{})

resp, err := g.Type(ctx, &proto.Empty{})
if test.expectErr && err == nil {
t.Fatalf("err expected, got nil")
}
Expand Down Expand Up @@ -539,14 +516,9 @@ func TestGRPCServer_Close(t *testing.T) {

for name, test := range tests {
t.Run(name, func(t *testing.T) {
g := gRPCServer{
impl: test.db,
}

// Context doesn't need to timeout since this is just passed through
ctx := context.Background()
idCtx, g := testGrpcServer(t, test.db)
_, err := g.Close(idCtx, &proto.Empty{})

_, err := g.Close(ctx, &proto.Empty{})
if test.expectErr && err == nil {
t.Fatalf("err expected, got nil")
}
Expand All @@ -562,6 +534,86 @@ func TestGRPCServer_Close(t *testing.T) {
}
}

func TestGetMultiplexIDFromContext(t *testing.T) {
type testCase struct {
ctx context.Context
expectedResp string
expectedErr error
}

tests := map[string]testCase{
"missing plugin multiplexing metadata": {
ctx: context.Background(),
expectedResp: "",
expectedErr: fmt.Errorf("missing plugin multiplexing metadata"),
},
"unexpected number of IDs in metadata": {
ctx: idCtx(t, "12345", "67891"),
expectedResp: "",
expectedErr: fmt.Errorf("unexpected number of IDs in metadata: (2)"),
},
"empty multiplex ID in metadata": {
ctx: idCtx(t, ""),
expectedResp: "",
expectedErr: fmt.Errorf("empty multiplex ID in metadata"),
},
"happy path, id is returned from metadata": {
ctx: idCtx(t, "12345"),
expectedResp: "12345",
expectedErr: nil,
},
}

for name, test := range tests {
t.Run(name, func(t *testing.T) {
resp, err := getMultiplexIDFromContext(test.ctx)

if test.expectedErr != nil && test.expectedErr.Error() != "" && err == nil {
t.Fatalf("err expected, got nil")
} else if !reflect.DeepEqual(err, test.expectedErr) {
t.Fatalf("Actual error: %#v\nExpected error: %#v", err, test.expectedErr)
}

if test.expectedErr != nil && test.expectedErr.Error() == "" && err != nil {
t.Fatalf("no error expected, got: %s", err)
}

if !reflect.DeepEqual(resp, test.expectedResp) {
t.Fatalf("Actual response: %#v\nExpected response: %#v", resp, test.expectedResp)
}
})
}
}

func testGrpcServer(t *testing.T, db Database) (context.Context, gRPCServer) {
t.Helper()
g := gRPCServer{
factoryFunc: func() (interface{}, error) {
return db, nil
},
instances: make(map[string]Database),
}

id := "12345"
idCtx := idCtx(t, id)
g.instances[id] = db

return idCtx, g
}

// idCtx is a test helper that will return a context with the IDs set in its
// metadata
func idCtx(t *testing.T, ids ...string) context.Context {
t.Helper()
// Context doesn't need to timeout since this is just passed through
ctx := context.Background()
md := metadata.MD{}
for _, id := range ids {
md.Append(pluginutil.MultiplexingCtxKey, id)
}
return metadata.NewIncomingContext(ctx, md)
}

func marshal(t *testing.T, m map[string]interface{}) *structpb.Struct {
t.Helper()

Expand Down
101 changes: 57 additions & 44 deletions sdk/helper/pluginutil/run_config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,19 +38,21 @@ func TestMakeConfig(t *testing.T) {
args: []string{"foo", "bar"},
sha256: []byte("some_sha256"),
env: []string{"initial=true"},
pluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
PluginClientConfig: PluginClientConfig{
PluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
},
},
HandshakeConfig: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
Logger: hclog.NewNullLogger(),
IsMetadataMode: true,
AutoMTLS: false,
},
hs: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
logger: hclog.NewNullLogger(),
isMetadataMode: true,
autoMTLS: false,
},

responseWrapInfoTimes: 0,
Expand Down Expand Up @@ -97,19 +99,21 @@ func TestMakeConfig(t *testing.T) {
args: []string{"foo", "bar"},
sha256: []byte("some_sha256"),
env: []string{"initial=true"},
pluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
PluginClientConfig: PluginClientConfig{
PluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
},
},
HandshakeConfig: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
Logger: hclog.NewNullLogger(),
IsMetadataMode: false,
AutoMTLS: false,
},
hs: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
logger: hclog.NewNullLogger(),
isMetadataMode: false,
autoMTLS: false,
},

responseWrapInfo: &wrapping.ResponseWrapInfo{
Expand Down Expand Up @@ -161,19 +165,21 @@ func TestMakeConfig(t *testing.T) {
args: []string{"foo", "bar"},
sha256: []byte("some_sha256"),
env: []string{"initial=true"},
pluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
PluginClientConfig: PluginClientConfig{
PluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
},
},
HandshakeConfig: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
Logger: hclog.NewNullLogger(),
IsMetadataMode: true,
AutoMTLS: true,
},
hs: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
logger: hclog.NewNullLogger(),
isMetadataMode: true,
autoMTLS: true,
},

responseWrapInfoTimes: 0,
Expand Down Expand Up @@ -220,19 +226,21 @@ func TestMakeConfig(t *testing.T) {
args: []string{"foo", "bar"},
sha256: []byte("some_sha256"),
env: []string{"initial=true"},
pluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
PluginClientConfig: PluginClientConfig{
PluginSets: map[int]plugin.PluginSet{
1: {
"bogus": nil,
},
},
HandshakeConfig: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
Logger: hclog.NewNullLogger(),
IsMetadataMode: false,
AutoMTLS: true,
},
hs: plugin.HandshakeConfig{
ProtocolVersion: 1,
MagicCookieKey: "magic_cookie_key",
MagicCookieValue: "magic_cookie_value",
},
logger: hclog.NewNullLogger(),
isMetadataMode: false,
autoMTLS: true,
},

responseWrapInfoTimes: 0,
Expand Down Expand Up @@ -329,6 +337,11 @@ type mockRunnerUtil struct {
mock.Mock
}

func (m *mockRunnerUtil) NewPluginClient(ctx context.Context, config PluginClientConfig) (PluginClient, error) {
args := m.Called(ctx, config)
return args.Get(0).(PluginClient), args.Error(1)
}

func (m *mockRunnerUtil) ResponseWrapData(ctx context.Context, data map[string]interface{}, ttl time.Duration, jwt bool) (*wrapping.ResponseWrapInfo, error) {
args := m.Called(ctx, data, ttl, jwt)
return args.Get(0).(*wrapping.ResponseWrapInfo), args.Error(1)
Expand Down

0 comments on commit 25620a1

Please sign in to comment.