diff --git a/storage/protobuf_store_test.go b/storage/protobuf_store_test.go index ef874ac081..d662ef536f 100644 --- a/storage/protobuf_store_test.go +++ b/storage/protobuf_store_test.go @@ -5,13 +5,17 @@ import ( "fmt" "io" "math/rand" + "net/http" + "net/http/httptest" "testing" - "github.com/flyteorg/flytestdlib/promutils" - + "github.com/flyteorg/stow/s3" "github.com/golang/protobuf/proto" errs "github.com/pkg/errors" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/flyteorg/flytestdlib/promutils" ) type mockProtoMessage struct { @@ -57,6 +61,36 @@ func TestDefaultProtobufStore(t *testing.T) { assert.Equal(t, int64(5), m.X) }) + t.Run("RefreshConfig", func(t *testing.T) { + testScope := promutils.NewTestScope() + s, err := NewDataStore(&Config{Type: TypeMemory}, testScope) + require.NoError(t, err) + require.IsType(t, DefaultProtobufStore{}, s.ComposedProtobufStore) + require.IsType(t, &InMemoryStore{}, s.ComposedProtobufStore.(DefaultProtobufStore).RawStore) + + oldMetrics := s.metrics + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + err = s.RefreshConfig(&Config{ + Type: TypeMinio, + Stow: StowConfig{ + Kind: TypeS3, + Config: map[string]string{ + s3.ConfigAccessKeyID: "key", + s3.ConfigSecretKey: "sec", + s3.ConfigEndpoint: server.URL, + }}, + InitContainer: "b"}) + + assert.NoError(t, err) + require.IsType(t, DefaultProtobufStore{}, s.ComposedProtobufStore) + assert.IsType(t, &StowStore{}, s.ComposedProtobufStore.(DefaultProtobufStore).RawStore) + assert.Equal(t, oldMetrics, s.metrics) + }) + t.Run("invalid type", func(t *testing.T) { testScope := promutils.NewTestScope() diff --git a/storage/rawstores.go b/storage/rawstores.go index 0e163d345a..f99cd8e3d3 100644 --- a/storage/rawstores.go +++ b/storage/rawstores.go @@ -108,6 +108,7 @@ func (ds *DataStore) RefreshConfig(cfg *Config) error { rawStore = newCachedRawStore(cfg, rawStore, ds.metrics.cacheMetrics) protoStore := NewDefaultProtobufStore(rawStore, ds.metrics.protoMetrics) newDS := NewCompositeDataStore(NewURLPathConstructor(), protoStore) + newDS.metrics = ds.metrics *ds = *newDS return nil }