Add unit tests for role_store and fix Remove

Signed-off-by: Mikkel Oscar Lyderik Larsen <mikkel.larsen@zalando.de>

go.mod

@@ -32,7 +32,7 @@ require (
 	github.com/smartystreets/assertions v0.0.0-20180820201707-7c9eb446e3cf // indirect
 	github.com/smartystreets/goconvey v0.0.0-20180222194500-ef6db91d284a // indirect
 	github.com/spf13/pflag v1.0.0 // indirect
-	github.com/stretchr/testify v1.2.2 // indirect
+	github.com/stretchr/testify v1.2.2
 	golang.org/x/crypto v0.0.0-20180330210355-12892e8c234f // indirect
 	golang.org/x/net v0.0.0-20180330215511-b68f30494add // indirect
 	golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f // indirect

role_store.go

@@ -61,8 +61,12 @@ func (s *RoleStore) Remove(role, namespace, name string) {
 
 	if ns, ok := s.Store[role]; ok {
 		if pods, ok := ns[namespace]; ok {
-			if _, ok := pods[name]; ok && len(pods) == 1 {
-				delete(ns, namespace)
+			if _, ok := pods[name]; ok {
+				if len(pods) == 1 {
+					delete(ns, namespace)
+				} else {
+					delete(pods, name)
+				}
 			}
 
 			if len(ns) == 0 {

role_store_test.go

@@ -0,0 +1,33 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRoleStore(t *testing.T) {
+	roleStore := NewRoleStore()
+	require.False(t, roleStore.Exists("role", "namespace"))
+
+	// Add and check
+	roleStore.Add("role", "namespace", "pod_name")
+	require.True(t, roleStore.Exists("role", "namespace"))
+
+	roleStore.Add("role", "namespace2", "pod_name")
+	require.True(t, roleStore.Exists("role", "namespace2"))
+
+	roleStore.Add("role", "namespace", "pod_name2")
+	require.True(t, roleStore.Exists("role", "namespace"))
+
+	// Remove and check
+	roleStore.Remove("role", "namespace", "pod_name")
+	require.True(t, roleStore.Exists("role", "namespace"))
+
+	roleStore.Remove("role", "namespace", "pod_name2")
+	require.False(t, roleStore.Exists("role", "namespace"))
+	require.True(t, roleStore.Exists("role", "namespace2"))
+
+	roleStore.Remove("role", "namespace2", "pod_name")
+	require.False(t, roleStore.Exists("role", "namespace2"))
+}