Refactor to allow easier testing

grafana · Sep 21, 2023 · 5d8997f · 5d8997f
1 parent 6723881
commit 5d8997f
Show file tree

Hide file tree

Showing 4 changed files with 167 additions and 69 deletions.
diff --git a/pkg/mimir/modules.go b/pkg/mimir/modules.go
@@ -172,12 +172,14 @@ func (t *Mimir) initActivityTracker() (services.Service, error) {
 	}), nil
 }
 
+var NewVault = vault.NewVault
+
 func (t *Mimir) initVault() (services.Service, error) {
 	if !t.Cfg.Vault.Enabled {
 		return nil, nil
 	}
 
-	v, err := vault.NewVault(t.Cfg.Vault)
+	v, err := NewVault(t.Cfg.Vault)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/mimir/modules_test.go b/pkg/mimir/modules_test.go
@@ -14,7 +14,10 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/server"
+	"github.com/grafana/mimir/pkg/util/test"
+	"github.com/grafana/mimir/pkg/vault"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -242,3 +245,83 @@ func TestMultiKVSetup(t *testing.T) {
 		})
 	}
 }
+
+func TestInitVault(t *testing.T) {
+	cfg := Config{
+		Server: server.Config{
+			HTTPTLSConfig: server.TLSConfig{
+				TLSCertPath: "test/secret1",
+				TLSKeyPath:  "test/secret2",
+				ClientCAs:   "test/secret1",
+			},
+			GRPCTLSConfig: server.TLSConfig{
+				TLSCertPath: "test/secret1",
+				TLSKeyPath:  "test/secret2",
+				ClientCAs:   "test/secret1",
+			},
+		},
+		Vault: vault.Config{
+			Enabled: true,
+		},
+	}
+
+	mimir := &Mimir{
+		Server: &server.Server{Registerer: prometheus.NewPedanticRegistry()},
+		Cfg:    cfg,
+	}
+
+	oldFunc := NewVault
+	defer func() {
+		NewVault = oldFunc
+	}()
+	NewVault = test.NewMockVault
+
+	_, err := mimir.initVault()
+	require.NoError(t, err)
+
+	// Check KVStore
+	require.NotNil(t, mimir.Cfg.MemberlistKV.TCPTransport.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Distributor.HATrackerConfig.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Alertmanager.ShardingRing.Common.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Compactor.ShardingRing.Common.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Distributor.DistributorRing.Common.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Ingester.IngesterRing.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Ruler.Ring.Common.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.StoreGateway.ShardingRing.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.QueryScheduler.ServiceDiscovery.SchedulerRing.KVStore.StoreConfig.Etcd.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.OverridesExporter.Ring.Common.KVStore.StoreConfig.Etcd.TLS.Reader)
+
+	// Check Redis Clients
+	require.NotNil(t, mimir.Cfg.BlocksStorage.BucketStore.IndexCache.BackendConfig.Redis.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.BlocksStorage.BucketStore.ChunksCache.BackendConfig.Redis.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.BlocksStorage.BucketStore.MetadataCache.BackendConfig.Redis.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Frontend.QueryMiddleware.ResultsCacheConfig.BackendConfig.Redis.TLS.Reader)
+
+	// Check GRPC Clients
+	require.NotNil(t, mimir.Cfg.IngesterClient.GRPCClientConfig.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Worker.GRPCClientConfig.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Querier.StoreGatewayClient.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Frontend.FrontendV2.GRPCClientConfig.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Ruler.ClientTLSConfig.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Ruler.Notifier.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.Alertmanager.AlertmanagerClient.GRPCClientConfig.TLS.Reader)
+	require.NotNil(t, mimir.Cfg.QueryScheduler.GRPCClientConfig.TLS.Reader)
+
+	// Check Server
+	require.Empty(t, mimir.Cfg.Server.HTTPTLSConfig.TLSCertPath)
+	require.Empty(t, mimir.Cfg.Server.HTTPTLSConfig.TLSKeyPath)
+	require.Empty(t, mimir.Cfg.Server.HTTPTLSConfig.ClientCAs)
+
+	require.Empty(t, mimir.Cfg.Server.GRPCTLSConfig.TLSCertPath)
+	require.Empty(t, mimir.Cfg.Server.GRPCTLSConfig.TLSKeyPath)
+	require.Empty(t, mimir.Cfg.Server.HTTPTLSConfig.ClientCAs)
+
+	require.Equal(t, "foo1", mimir.Cfg.Server.HTTPTLSConfig.TLSCert)
+	require.Equal(t, config.Secret("foo2"), mimir.Cfg.Server.HTTPTLSConfig.TLSKey)
+	require.Equal(t, "foo1", mimir.Cfg.Server.HTTPTLSConfig.ClientCAsText)
+
+	require.Equal(t, "foo1", mimir.Cfg.Server.GRPCTLSConfig.TLSCert)
+	require.Equal(t, config.Secret("foo2"), mimir.Cfg.Server.GRPCTLSConfig.TLSKey)
+	require.Equal(t, "foo1", mimir.Cfg.Server.GRPCTLSConfig.ClientCAsText)
+
+}
diff --git a/pkg/util/test/vault.go b/pkg/util/test/vault.go
@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package test
+
+import (
+	"context"
+
+	"github.com/grafana/mimir/pkg/vault"
+	hashivault "github.com/hashicorp/vault/api"
+	"github.com/pkg/errors"
+)
+
+type mockKVStore struct {
+	values map[string]mockValue
+}
+
+type mockValue struct {
+	secret *hashivault.KVSecret
+	err    error
+}
+
+func newMockKVStore() *mockKVStore {
+	return &mockKVStore{
+		values: map[string]mockValue{
+			"test/secret1": {
+				secret: &hashivault.KVSecret{
+					Data: map[string]interface{}{
+						"value": "foo1",
+					},
+				},
+				err: nil,
+			},
+			"test/secret2": {
+				secret: &hashivault.KVSecret{
+					Data: map[string]interface{}{
+						"value": "foo2",
+					},
+				},
+				err: nil,
+			},
+			"test/secret3": {
+				secret: nil,
+				err:    errors.New("non-existent path"),
+			},
+			"test/secret4": {
+				secret: nil,
+				err:    nil,
+			},
+			"test/secret5": {
+				secret: &hashivault.KVSecret{
+					Data: map[string]interface{}{
+						"value": 123,
+					},
+				},
+				err: nil,
+			},
+			"test/secret6": {
+				secret: &hashivault.KVSecret{
+					Data: map[string]interface{}{
+						"value": nil,
+					},
+				},
+				err: nil,
+			},
+		},
+	}
+}
+
+func (m *mockKVStore) Get(_ context.Context, path string) (*hashivault.KVSecret, error) {
+	return m.values[path].secret, m.values[path].err
+}
+
+func NewMockVault(vault.Config) (*vault.Vault, error) {
+	return &vault.Vault{
+		KVStore: newMockKVStore(),
+	}, nil
+}
diff --git a/pkg/vault/vault_test.go b/pkg/vault/vault_test.go
@@ -1,22 +1,18 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
-package vault
+package vault_test
 
 import (
-	"context"
-	"errors"
 	"testing"
 
-	hashivault "github.com/hashicorp/vault/api"
+	"github.com/grafana/mimir/pkg/util/test"
+	"github.com/grafana/mimir/pkg/vault"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestReadSecret(t *testing.T) {
-	mockKVStore := newMockKVStore()
-	mimirVaultClient := Vault{
-		KVStore: mockKVStore,
-	}
+	mimirVaultClient, _ := test.NewMockVault(vault.Config{})
 
 	tests := map[string]struct {
 		path          string
@@ -69,63 +65,3 @@ func TestReadSecret(t *testing.T) {
 		})
 	}
 }
-
-type mockKVStore struct {
-	values map[string]mockValue
-}
-
-type mockValue struct {
-	secret *hashivault.KVSecret
-	err    error
-}
-
-func newMockKVStore() *mockKVStore {
-	return &mockKVStore{
-		values: map[string]mockValue{
-			"test/secret1": {
-				secret: &hashivault.KVSecret{
-					Data: map[string]interface{}{
-						"value": "foo1",
-					},
-				},
-				err: nil,
-			},
-			"test/secret2": {
-				secret: &hashivault.KVSecret{
-					Data: map[string]interface{}{
-						"value": "foo2",
-					},
-				},
-				err: nil,
-			},
-			"test/secret3": {
-				secret: nil,
-				err:    errors.New("non-existent path"),
-			},
-			"test/secret4": {
-				secret: nil,
-				err:    nil,
-			},
-			"test/secret5": {
-				secret: &hashivault.KVSecret{
-					Data: map[string]interface{}{
-						"value": 123,
-					},
-				},
-				err: nil,
-			},
-			"test/secret6": {
-				secret: &hashivault.KVSecret{
-					Data: map[string]interface{}{
-						"value": nil,
-					},
-				},
-				err: nil,
-			},
-		},
-	}
-}
-
-func (m *mockKVStore) Get(_ context.Context, path string) (*hashivault.KVSecret, error) {
-	return m.values[path].secret, m.values[path].err
-}