Added Initializing auth client from config method (#140)

* Added Initializing auth client from config method

Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>

* Using sync.Once only once for initializing all the clients

Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>

* Added mock generation for AuthServiceClient

Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>

* Added clientset wrapping adminClient and authClient

Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>

* Added initialization of the mock client set

Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>

* Exposed the AdminServiceClient and AuthServiceClient for mocking

Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>

* Incorporated the feedback

Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>

Co-authored-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>
Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

clients/go/admin/client.go

@@ -13,6 +13,7 @@ import (
 	"github.com/flyteorg/flyteidl/clients/go/admin/mocks"
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
 	"github.com/flyteorg/flytestdlib/logger"
+
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
 	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
@@ -26,11 +27,32 @@ var (
 	adminConnection *grpc.ClientConn
 )
 
+// Clientset contains the clients exposed to communicate with various admin services.
+type Clientset struct {
+	AdminServiceClient service.AdminServiceClient
+	AuthServiceClient  service.AuthServiceClient
+}
+
+// AdminClient retrieves the AdminServiceClient
+func (c Clientset) AdminClient() service.AdminServiceClient {
+	return c.AdminServiceClient
+}
+
+// AuthClient retrieves the AuthServiceClient
+func (c Clientset) AuthClient() service.AuthServiceClient {
+	return c.AuthServiceClient
+}
+
 func NewAdminClient(ctx context.Context, conn *grpc.ClientConn) service.AdminServiceClient {
 	logger.Infof(ctx, "Initialized Admin client")
 	return service.NewAdminServiceClient(conn)
 }
 
+func NewAuthClient(ctx context.Context, conn *grpc.ClientConn) service.AuthServiceClient {
+	logger.Infof(ctx, "Initialized Auth client")
+	return service.NewAuthServiceClient(conn)
+}
+
 func GetAdditionalAdminClientConfigOptions(cfg Config) []grpc.DialOption {
 	opts := make([]grpc.DialOption, 0, 2)
 	backoffConfig := grpc.BackoffConfig{
@@ -144,6 +166,21 @@ func InitializeAdminClient(ctx context.Context, cfg Config) service.AdminService
 	return NewAdminClient(ctx, adminConnection)
 }
 
+// Create an AdminClient and AuthServiceClient with a shared Admin connection for the process
+func InitializeClients(ctx context.Context, cfg Config) (*Clientset, error) {
+	once.Do(func() {
+		var err error
+		adminConnection, err = NewAdminConnection(ctx, cfg)
+		if err != nil {
+			logger.Panicf(ctx, "failed to initialize Admin connection. Err: %s", err.Error())
+		}
+	})
+	var cs Clientset
+	cs.AdminServiceClient = NewAdminClient(ctx, adminConnection)
+	cs.AuthServiceClient = NewAuthClient(ctx, adminConnection)
+	return &cs, nil
+}
+
 func InitializeAdminClientFromConfig(ctx context.Context) (service.AdminServiceClient, error) {
 	cfg := GetConfig(ctx)
 	if cfg == nil {
@@ -152,7 +189,20 @@ func InitializeAdminClientFromConfig(ctx context.Context) (service.AdminServiceC
 	return InitializeAdminClient(ctx, *cfg), nil
 }
 
+func InitializeClientsFromConfig(ctx context.Context) (*Clientset, error) {
+	cfg := GetConfig(ctx)
+	if cfg == nil {
+		return nil, fmt.Errorf("retrieved Nil config for [%s] key", configSectionKey)
+	}
+	return InitializeClients(ctx, *cfg)
+}
+
 func InitializeMockAdminClient() service.AdminServiceClient {
 	logger.Infof(context.TODO(), "Initialized Mock Admin client")
 	return &mocks.AdminServiceClient{}
 }
+
+func InitializeMockClientset() *Clientset {
+	logger.Infof(context.TODO(), "Initialized Mock Clientset")
+	return &Clientset{AdminServiceClient: &mocks.AdminServiceClient{}, AuthServiceClient: &mocks.AuthServiceClient{}}
+}

clients/go/admin/client_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/flyteorg/flytestdlib/config"
+
 	"github.com/stretchr/testify/assert"
 )
 
@@ -22,6 +23,12 @@ func TestInitializeAndGetAdminClient(t *testing.T) {
 		}))
 	})
 
+	t.Run("legal-from-config", func(t *testing.T) {
+		authClient, err := InitializeAdminClientFromConfig(ctx)
+		assert.NoError(t, err)
+		assert.NotNil(t, authClient)
+	})
+
 	t.Run("illegal", func(t *testing.T) {
 		adminConnection = nil
 		once = sync.Once{}
@@ -34,6 +41,13 @@ func TestInitializeMockAdminClient(t *testing.T) {
 	assert.NotNil(t, c)
 }
 
+func TestInitializeMockClientset(t *testing.T) {
+	c := InitializeMockClientset()
+	assert.NotNil(t, c)
+	assert.NotNil(t, c.AdminServiceClient)
+	assert.NotNil(t, c.AuthServiceClient)
+}
+
 func TestGetAdditionalAdminClientConfigOptions(t *testing.T) {
 	u, _ := url.Parse("localhost:8089")
 	adminServiceConfig := Config{
@@ -45,3 +59,24 @@ func TestGetAdditionalAdminClientConfigOptions(t *testing.T) {
 	opts := GetAdditionalAdminClientConfigOptions(adminServiceConfig)
 	assert.Equal(t, 2, len(opts))
 }
+
+func TestInitializeClients(t *testing.T) {
+	ctx := context.TODO()
+	t.Run("legal", func(t *testing.T) {
+		u, err := url.Parse("http://localhost:8089")
+		assert.NoError(t, err)
+		clientSet, err := InitializeClients(ctx, Config{Endpoint: config.URL{URL: *u}})
+		assert.NotNil(t, clientSet)
+		assert.NotNil(t, clientSet.AdminClient())
+		assert.NotNil(t, clientSet.AuthClient())
+		assert.Nil(t, err)
+	})
+
+	t.Run("legal-from-config", func(t *testing.T) {
+		clientSet, err := InitializeClientsFromConfig(ctx)
+		assert.NoError(t, err)
+		assert.NotNil(t, clientSet)
+		assert.NotNil(t, clientSet.AuthClient())
+		assert.NotNil(t, clientSet.AdminClient())
+	})
+}

generate_mocks.sh

@@ -3,4 +3,5 @@ set -e
 set -x
 
 mockery -dir=gen/pb-go/flyteidl/service/ -name=AdminServiceClient -output=clients/go/admin/mocks
+mockery -dir=gen/pb-go/flyteidl/service/ -name=AuthServiceClient -output=clients/go/admin/mocks
 mockery -dir=gen/pb-go/flyteidl/datacatalog/ -name=DataCatalogClient -output=clients/go/datacatalog/mocks