plugins: add unit tests for close client methods

Signed-off-by: Steve Kriss <steve@heptio.com>

pkg/plugin/client_builder.go

@@ -38,6 +38,6 @@ func (b *clientBuilder) withCommand(name string, args ...string) *clientBuilder
 	return b
 }
 
-func (b *clientBuilder) client() *hcplugin.Client {
+func (b *clientBuilder) client() pluginClient {
 	return hcplugin.NewClient(b.config)
 }

pkg/plugin/client_store.go

@@ -3,7 +3,6 @@ package plugin
 import (
 	"sync"
 
-	plugin "github.com/hashicorp/go-plugin"
 	"github.com/pkg/errors"
 )
 
@@ -20,7 +19,7 @@ type clientKey struct {
 
 func newClientStore() *clientStore {
 	return &clientStore{
-		clients: make(map[clientKey]map[string]*plugin.Client),
+		clients: make(map[clientKey]map[string]pluginClient),
 		lock:    &sync.RWMutex{},
 	}
 }
@@ -33,13 +32,13 @@ type clientStore struct {
 	// kind and scope (e.g. all BackupItemActions for a given
 	// backup), and efficient lookup by kind+name+scope (e.g.
 	// the AWS ObjectStore.)
-	clients map[clientKey]map[string]*plugin.Client
+	clients map[clientKey]map[string]pluginClient
 	lock    *sync.RWMutex
 }
 
 // get returns a plugin client for the given kind/name/scope, or an error if none
 // is found.
-func (s *clientStore) get(kind PluginKind, name, scope string) (*plugin.Client, error) {
+func (s *clientStore) get(kind PluginKind, name, scope string) (pluginClient, error) {
 	s.lock.RLock()
 	defer s.lock.RUnlock()
 
@@ -54,12 +53,12 @@ func (s *clientStore) get(kind PluginKind, name, scope string) (*plugin.Client,
 
 // list returns all plugin clients for the given kind/scope, or an
 // error if none are found.
-func (s *clientStore) list(kind PluginKind, scope string) ([]*plugin.Client, error) {
+func (s *clientStore) list(kind PluginKind, scope string) ([]pluginClient, error) {
 	s.lock.RLock()
 	defer s.lock.RUnlock()
 
 	if forScope, found := s.clients[clientKey{kind, scope}]; found {
-		var clients []*plugin.Client
+		var clients []pluginClient
 
 		for _, client := range forScope {
 			clients = append(clients, client)
@@ -73,11 +72,11 @@ func (s *clientStore) list(kind PluginKind, scope string) ([]*plugin.Client, err
 
 // listAll returns all plugin clients for all kinds/scopes, or a
 // zero-valued slice if there are none.
-func (s *clientStore) listAll() []*plugin.Client {
+func (s *clientStore) listAll() []pluginClient {
 	s.lock.RLock()
 	defer s.lock.RUnlock()
 
-	var clients []*plugin.Client
+	var clients []pluginClient
 	for _, pluginsByName := range s.clients {
 		for name := range pluginsByName {
 			clients = append(clients, pluginsByName[name])
@@ -88,14 +87,14 @@ func (s *clientStore) listAll() []*plugin.Client {
 }
 
 // add stores a plugin client for the given kind/name/scope.
-func (s *clientStore) add(client *plugin.Client, kind PluginKind, name, scope string) {
+func (s *clientStore) add(client pluginClient, kind PluginKind, name, scope string) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
 	key := clientKey{kind, scope}
 
 	if _, found := s.clients[key]; !found {
-		s.clients[key] = make(map[string]*plugin.Client)
+		s.clients[key] = make(map[string]pluginClient)
 	}
 
 	s.clients[key][name] = client
@@ -125,5 +124,5 @@ func (s *clientStore) clear() {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	s.clients = make(map[clientKey]map[string]*plugin.Client)
+	s.clients = make(map[clientKey]map[string]pluginClient)
 }

pkg/plugin/manager.go

@@ -164,7 +164,12 @@ func pluginForKind(kind PluginKind) plugin.Plugin {
 	}
 }
 
-func getPluginInstance(client *plugin.Client, kind PluginKind) (interface{}, error) {
+type pluginClient interface {
+	Client() (plugin.ClientProtocol, error)
+	Kill()
+}
+
+func getPluginInstance(client pluginClient, kind PluginKind) (interface{}, error) {
 	protocolClient, err := client.Client()
 	if err != nil {
 		return nil, errors.WithStack(err)

pkg/plugin/manager_test.go

@@ -0,0 +1,179 @@
+/*
+Copyright 2018 the Heptio Ark contributors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package plugin
+
+import (
+	"testing"
+
+	plugin "github.com/hashicorp/go-plugin"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type fakePluginClient struct {
+	terminated map[*fakePluginClient]bool
+}
+
+func (c *fakePluginClient) Kill() {
+	c.terminated[c] = true
+}
+
+func (c *fakePluginClient) Client() (plugin.ClientProtocol, error) {
+	panic("not implemented")
+}
+
+func TestCloseAllClients(t *testing.T) {
+	var (
+		store       = newClientStore()
+		terminated  = make(map[*fakePluginClient]bool)
+		clientCount = 0
+	)
+
+	for _, kind := range AllPluginKinds {
+		for _, scope := range []string{"scope-1", "scope-2"} {
+			for _, name := range []string{"name-1", "name-2"} {
+				client := &fakePluginClient{terminated: terminated}
+				terminated[client] = false
+				store.add(client, kind, name, scope)
+				clientCount++
+			}
+		}
+	}
+
+	// verify setup
+	require.Len(t, terminated, clientCount)
+	for _, status := range terminated {
+		require.False(t, status)
+	}
+
+	m := &manager{clientStore: store}
+	m.CloseAllClients()
+
+	// we should have no additions to or removals from the `terminated` map
+	assert.Len(t, terminated, clientCount)
+
+	// all clients should have their entry in the `terminated` map flipped to true
+	for _, status := range terminated {
+		assert.True(t, status)
+	}
+	// the store's `clients` map should be empty
+	assert.Len(t, store.clients, 0)
+}
+
+func TestCloseBackupItemActions(t *testing.T) {
+	var (
+		store                = newClientStore()
+		terminated           = make(map[*fakePluginClient]bool)
+		clientCount          = 0
+		expectedTerminations = make(map[*fakePluginClient]bool)
+		backupName           = "backup-1"
+	)
+
+	for _, kind := range AllPluginKinds {
+		for _, scope := range []string{"backup-1", "backup-2"} {
+			for _, name := range []string{"name-1", "name-2"} {
+				client := &fakePluginClient{terminated: terminated}
+				terminated[client] = false
+				store.add(client, kind, name, scope)
+				clientCount++
+
+				if kind == PluginKindBackupItemAction && scope == backupName {
+					expectedTerminations[client] = true
+				}
+			}
+		}
+	}
+
+	// verify setup
+	require.Len(t, terminated, clientCount)
+	for _, status := range terminated {
+		require.False(t, status)
+	}
+
+	m := &manager{clientStore: store}
+	m.CloseBackupItemActions(backupName)
+
+	// we should have no additions to or removals from the `terminated` map
+	assert.Len(t, terminated, clientCount)
+
+	// only those clients that we expected to be terminated should have
+	// their entry in the `terminated` map flipped to true
+	for client, status := range terminated {
+		_, ok := expectedTerminations[client]
+		assert.Equal(t, ok, status)
+	}
+
+	// clients for the kind/scope should have been removed
+	_, err := store.list(PluginKindBackupItemAction, backupName)
+	assert.EqualError(t, err, "clients not found")
+
+	// total number of clients should decrease by the number of terminated
+	// clients
+	assert.Len(t, store.listAll(), clientCount-len(expectedTerminations))
+}
+
+func TestCloseRestoreItemActions(t *testing.T) {
+	var (
+		store                = newClientStore()
+		terminated           = make(map[*fakePluginClient]bool)
+		clientCount          = 0
+		expectedTerminations = make(map[*fakePluginClient]bool)
+		restoreName          = "restore-2"
+	)
+
+	for _, kind := range AllPluginKinds {
+		for _, scope := range []string{"restore-1", "restore-2"} {
+			for _, name := range []string{"name-1", "name-2"} {
+				client := &fakePluginClient{terminated: terminated}
+				terminated[client] = false
+				store.add(client, kind, name, scope)
+				clientCount++
+
+				if kind == PluginKindRestoreItemAction && scope == restoreName {
+					expectedTerminations[client] = true
+				}
+			}
+		}
+	}
+
+	// verify setup
+	require.Len(t, terminated, clientCount)
+	for _, status := range terminated {
+		require.False(t, status)
+	}
+
+	m := &manager{clientStore: store}
+	m.CloseRestoreItemActions(restoreName)
+
+	// we should have no additions to or removals from the `terminated` map
+	assert.Len(t, terminated, clientCount)
+
+	// only those clients that we expected to be terminated should have
+	// their entry in the `terminated` map flipped to true
+	for client, status := range terminated {
+		_, ok := expectedTerminations[client]
+		assert.Equal(t, ok, status)
+	}
+
+	// clients for the kind/scope should have been removed
+	_, err := store.list(PluginKindRestoreItemAction, restoreName)
+	assert.EqualError(t, err, "clients not found")
+
+	// total number of clients should decrease by the number of terminated
+	// clients
+	assert.Len(t, store.listAll(), clientCount-len(expectedTerminations))
+}