more tests

Signed-off-by: Kevin Su <[email protected]>

flyteplugins/go/tasks/pluginmachinery/core/secret_manager.go

@@ -2,8 +2,6 @@ package core
 
 import "context"
 
-//go:generate mockery -all -output=./mocks -case=underscore
-
 type SecretManager interface {
 	Get(ctx context.Context, key string) (string, error)
 }

flyteplugins/go/tasks/plugins/webapi/agent/integration_test.go

@@ -104,6 +104,8 @@ func TestEndToEnd(t *testing.T) {
 				},
 			},
 		}
+		template.SecurityContext = &flyteIdlCore.SecurityContext{Connection: "openai"}
+
 		expectedOutputs, err := coreutils.MakeLiteralMap(map[string]interface{}{"x": 1})
 		assert.NoError(t, err)
 		phase := tests.RunPluginEndToEndTest(t, plugin, &template, inputs, expectedOutputs, nil, iter)

flyteplugins/tests/end_to_end.go

@@ -231,6 +231,10 @@ func RunPluginEndToEndTest(t *testing.T, executor pluginCore.Plugin, template *i
 	secretManager := &coreMocks.SecretManager{}
 	secretManager.OnGet(ctx, mock.Anything).Return("fake-token", nil)
 
+	connection := idlCore.Connection{Secrets: map[string]string{"OPENAI_API_KEY": "123"}}
+	connectionManager := &coreMocks.ConnectionManager{}
+	connectionManager.OnGet(ctx, mock.Anything).Return(connection, nil)
+
 	tCtx := &coreMocks.TaskExecutionContext{}
 	tCtx.OnInputReader().Return(inputReader)
 	tCtx.OnTaskRefreshIndicator().Return(func(ctx context.Context) {})
@@ -245,6 +249,7 @@ func RunPluginEndToEndTest(t *testing.T, executor pluginCore.Plugin, template *i
 	tCtx.OnResourceManager().Return(resourceManager)
 	tCtx.OnMaxDatasetSizeBytes().Return(1000000)
 	tCtx.OnSecretManager().Return(secretManager)
+	tCtx.OnConnectionManager().Return(connectionManager)
 
 	trns := pluginCore.DoTransition(pluginCore.PhaseInfoQueued(time.Now(), 0, ""))
 	for !trns.Info().Phase().IsTerminal() {

flytepropeller/pkg/controller/nodes/task/connectionmanager/config.go

@@ -19,6 +19,10 @@ type Config struct {
 	Connection map[string]Connection `json:"connection" pflag:", the connection that saves the secrets and configs"`
 }
 
+func SetConfig(cfg *Config) error {
+	return section.SetConfig(cfg)
+}
+
 func GetConfig() *Config {
 	return section.GetConfig().(*Config)
 }

flytepropeller/pkg/controller/nodes/task/connectionmanager/connection.go

@@ -3,38 +3,39 @@ package connectionmanager
 import (
 	"context"
 	"fmt"
-	"os"
-
 	flyteidl "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
+	pluginCore "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"
+	"github.com/flyteorg/flyte/flytepropeller/pkg/controller/nodes/task/secretmanager"
+	"github.com/flyteorg/flyte/flytestdlib/logger"
 )
 
 // FileEnvConnectionManager allows retrieving secrets mounted to this process through Env Vars or Files.
-type FileEnvConnectionManager struct{}
+type FileEnvConnectionManager struct {
+	secretManager pluginCore.SecretManager
+}
 
 // Get retrieves a secret from the environment of the running process or from a file.
-func (f FileEnvConnectionManager) Get(_ context.Context, key string) (flyteidl.Connection, error) {
+func (f FileEnvConnectionManager) Get(ctx context.Context, key string) (flyteidl.Connection, error) {
 	cfg := GetConfig()
 	connection, ok := cfg.Connection[key]
 	if !ok {
 		return flyteidl.Connection{}, fmt.Errorf("connection not found: [%s]", key)
 	}
 	secret := make(map[string]string)
 	for k, v := range connection.Secrets {
-		// TODO: Read the secret from a local file
-		val, ok := os.LookupEnv(v)
-		if !ok {
-			return flyteidl.Connection{}, fmt.Errorf("secret not found in env [%s]", v)
+		val, err := f.secretManager.Get(ctx, v)
+		if err != nil {
+			logger.Errorf(ctx, "failed to get secret [%s] for connection [%s] with error: %v", v, k, err)
+			return flyteidl.Connection{}, err
 		}
 		secret[k] = val
 	}
-	config := make(map[string]string)
-	for k, v := range connection.Configs {
-		config[k] = v
-	}
 
-	return flyteidl.Connection{Secrets: secret, Configs: config}, nil
+	return flyteidl.Connection{Secrets: secret, Configs: connection.Configs}, nil
 }
 
 func NewFileEnvConnectionManager() FileEnvConnectionManager {
-	return FileEnvConnectionManager{}
+	return FileEnvConnectionManager{
+		secretManager: secretmanager.NewFileEnvSecretManager(secretmanager.GetConfig()),
+	}
 }

flytepropeller/pkg/controller/nodes/task/connectionmanager/connection_test.go

@@ -0,0 +1,38 @@
+package connectionmanager
+
+import (
+	"context"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"testing"
+
+	coreMocks "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core/mocks"
+)
+
+func TestConnectionManager(t *testing.T) {
+	ctx := context.Background()
+	fakeSecretManager := &coreMocks.SecretManager{}
+	fakeSecretManager.OnGet(ctx, mock.Anything).Return("fake-token", nil)
+
+	connectionManager := NewFileEnvConnectionManager()
+	connectionManager.secretManager = fakeSecretManager
+
+	cfg := defaultConfig
+	cfg.Connection = map[string]Connection{
+		"openai": {
+			Secrets: map[string]string{
+				"openai_api_key": "api_key",
+			},
+			Configs: map[string]string{
+				"openai_organization": "flyteorg",
+			},
+		},
+	}
+	err := SetConfig(cfg)
+	assert.Nil(t, err)
+
+	connection, err := connectionManager.Get(ctx, "openai")
+	assert.Nil(t, err)
+	assert.Equal(t, "fake-token", connection.Secrets["openai_api_key"])
+	assert.Equal(t, "flyteorg", connection.Configs["openai_organization"])
+}