Added non zero return code on error and AlreadyExists as success cond…

…ition (flyteorg#71) * Added non zero return code on error and also made AlreadyExists as success condition Signed-off-by: Prafulla Mahindrakar <[email protected]> * Added coverage Signed-off-by: Prafulla Mahindrakar <[email protected]>
robert-ulbrich-mercedes-benz · Jul 2, 2024 · a3599fc · a3599fc
1 parent 5e7db6d
commit a3599fc
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 15 deletions.
diff --git a/flytectl/cmd/register/files.go b/flytectl/cmd/register/files.go
@@ -94,5 +94,5 @@ func registerFromFilesFunc(ctx context.Context, args []string, cmdCtx cmdCore.Co
 			logger.Errorf(ctx, "unable to delete temp dir %v due to %v", tmpDir, _err)
 		}
 	}
-	return nil
+	return _err
 }
diff --git a/flytectl/cmd/register/register_util.go b/flytectl/cmd/register/register_util.go
@@ -23,6 +23,8 @@ import (
 
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/golang/protobuf/proto"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 const registrationProjectPattern = "{{ registration.project }}"
@@ -40,10 +42,10 @@ type HTTPClient interface {
 	Do(req *http.Request) (*http.Response, error)
 }
 
-var Client HTTPClient
+var httpClient HTTPClient
 
 func init() {
-	Client = &http.Client{}
+	httpClient = &http.Client{}
 }
 
 var projectColumns = []printer.Column{
@@ -214,7 +216,7 @@ func DownloadFileFromHTTP(ctx context.Context, ref storage.DataReference) (io.Re
 	if err != nil {
 		return nil, err
 	}
-	resp, err := Client.Do(req)
+	resp, err := httpClient.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -320,7 +322,13 @@ func registerFile(ctx context.Context, fileName string, registerResults []Result
 	}
 	logger.Debugf(ctx, "Hydrated spec : %v", getJSONSpec(spec))
 	if err := register(ctx, spec, cmdCtx); err != nil {
-		registerResult = Result{Name: fileName, Status: "Failed", Info: fmt.Sprintf("Error registering file due to %v", err)}
+		// If error is AlreadyExists then dont consider this to be an error but just a warning state
+		if grpcError := status.Code(err); grpcError == codes.AlreadyExists {
+			registerResult = Result{Name: fileName, Status: "Success", Info: fmt.Sprintf("%v", grpcError.String())}
+			err = nil
+		} else {
+			registerResult = Result{Name: fileName, Status: "Failed", Info: fmt.Sprintf("Error registering file due to %v", err)}
+		}
 		registerResults = append(registerResults, registerResult)
 		return registerResults, err
 	}

diff --git a/flytectl/cmd/register/register_util_test.go b/flytectl/cmd/register/register_util_test.go
@@ -10,26 +10,36 @@ import (
 	"strings"
 	"testing"
 
+	cmdCore "github.com/flyteorg/flytectl/cmd/core"
+	u "github.com/flyteorg/flytectl/cmd/testutils"
+	"github.com/flyteorg/flyteidl/clients/go/admin/mocks"
+
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
-type MockClient struct {
+type MockHTTPClient struct {
 	DoFunc func(req *http.Request) (*http.Response, error)
 }
 
-func (m *MockClient) Do(req *http.Request) (*http.Response, error) {
+func (m *MockHTTPClient) Do(req *http.Request) (*http.Response, error) {
 	return GetDoFunc(req)
 }
 
 var (
-	ctx       context.Context
-	args      []string
-	GetDoFunc func(req *http.Request) (*http.Response, error)
+	ctx             context.Context
+	mockAdminClient *mocks.AdminServiceClient
+	cmdCtx          cmdCore.CommandContext
+	args            []string
+	GetDoFunc       func(req *http.Request) (*http.Response, error)
 )
 
-func setup() {
-	ctx = context.Background()
-	Client = &MockClient{}
+var setup = u.Setup
+
+func registerFilesSetup() {
+	httpClient = &MockHTTPClient{}
 	validTar, err := os.Open("testdata/valid-register.tar")
 	if err != nil {
 		fmt.Printf("unexpected error: %v", err)
@@ -41,10 +51,14 @@ func setup() {
 	GetDoFunc = func(*http.Request) (*http.Response, error) {
 		return response, nil
 	}
+	ctx = u.Ctx
+	mockAdminClient = u.MockClient
+	cmdCtx = cmdCore.NewCommandContext(mockAdminClient, u.MockOutStream)
 }
 
 func TestGetSortedFileList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = false
 	args = []string{"file2", "file1"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -56,6 +70,7 @@ func TestGetSortedFileList(t *testing.T) {
 
 func TestGetSortedArchivedFileWithParentFolderList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"testdata/valid-parent-folder-register.tar"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -72,6 +87,7 @@ func TestGetSortedArchivedFileWithParentFolderList(t *testing.T) {
 
 func TestGetSortedArchivedFileList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"testdata/valid-register.tar"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -88,6 +104,7 @@ func TestGetSortedArchivedFileList(t *testing.T) {
 
 func TestGetSortedArchivedFileUnorderedList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"testdata/valid-unordered-register.tar"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -104,6 +121,7 @@ func TestGetSortedArchivedFileUnorderedList(t *testing.T) {
 
 func TestGetSortedArchivedCorruptedFileList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"testdata/invalid.tar"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -116,6 +134,7 @@ func TestGetSortedArchivedCorruptedFileList(t *testing.T) {
 
 func TestGetSortedArchivedTgzList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"testdata/valid-register.tgz"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -144,6 +163,7 @@ func TestGetSortedArchivedCorruptedTgzFileList(t *testing.T) {
 
 func TestGetSortedArchivedInvalidArchiveFileList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"testdata/invalid-extension-register.zip"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -169,6 +189,7 @@ func TestGetSortedArchivedFileThroughInvalidHttpList(t *testing.T) {
 
 func TestGetSortedArchivedFileThroughValidHttpList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"http://dummyhost:80/testdata/valid-register.tar"}
 	fileList, tmpDir, err := getSortedFileList(ctx, args)
@@ -185,6 +206,7 @@ func TestGetSortedArchivedFileThroughValidHttpList(t *testing.T) {
 
 func TestGetSortedArchivedFileThroughValidHttpWithNullContextList(t *testing.T) {
 	setup()
+	registerFilesSetup()
 	filesConfig.Archive = true
 	args = []string{"http://dummyhost:80/testdata/valid-register.tar"}
 	ctx = nil
@@ -196,3 +218,64 @@ func TestGetSortedArchivedFileThroughValidHttpWithNullContextList(t *testing.T)
 	// Clean up the temp directory.
 	assert.Nil(t, os.RemoveAll(tmpDir), "unable to delete temp dir %v", tmpDir)
 }
+
+func TestRegisterFile(t *testing.T) {
+	t.Run("Successful run", func(t *testing.T) {
+		setup()
+		registerFilesSetup()
+		mockAdminClient.OnCreateTaskMatch(mock.Anything, mock.Anything).Return(nil, nil)
+		args = []string{"testdata/69_core.flyte_basics.lp.greet_1.pb"}
+		var registerResults []Result
+		results, err := registerFile(ctx, args[0], registerResults, cmdCtx)
+		assert.Equal(t, 1, len(results))
+		assert.Nil(t, err)
+	})
+	t.Run("Non existent file", func(t *testing.T) {
+		setup()
+		registerFilesSetup()
+		args = []string{"testdata/non-existent.pb"}
+		var registerResults []Result
+		results, err := registerFile(ctx, args[0], registerResults, cmdCtx)
+		assert.Equal(t, 1, len(results))
+		assert.Equal(t, "Failed", results[0].Status)
+		assert.Equal(t, "Error reading file due to open testdata/non-existent.pb: no such file or directory", results[0].Info)
+		assert.NotNil(t, err)
+	})
+	t.Run("unmarhal failure", func(t *testing.T) {
+		setup()
+		registerFilesSetup()
+		args = []string{"testdata/valid-register.tar"}
+		var registerResults []Result
+		results, err := registerFile(ctx, args[0], registerResults, cmdCtx)
+		assert.Equal(t, 1, len(results))
+		assert.Equal(t, "Failed", results[0].Status)
+		assert.Equal(t, "Error unmarshalling file due to failed unmarshalling file testdata/valid-register.tar", results[0].Info)
+		assert.NotNil(t, err)
+	})
+	t.Run("AlreadyExists", func(t *testing.T) {
+		setup()
+		registerFilesSetup()
+		mockAdminClient.OnCreateTaskMatch(mock.Anything, mock.Anything).Return(nil,
+			status.Error(codes.AlreadyExists, "AlreadyExists"))
+		args = []string{"testdata/69_core.flyte_basics.lp.greet_1.pb"}
+		var registerResults []Result
+		results, err := registerFile(ctx, args[0], registerResults, cmdCtx)
+		assert.Equal(t, 1, len(results))
+		assert.Equal(t, "Success", results[0].Status)
+		assert.Equal(t, "AlreadyExists", results[0].Info)
+		assert.Nil(t, err)
+	})
+	t.Run("Registration Error", func(t *testing.T) {
+		setup()
+		registerFilesSetup()
+		mockAdminClient.OnCreateTaskMatch(mock.Anything, mock.Anything).Return(nil,
+			status.Error(codes.InvalidArgument, "Invalid"))
+		args = []string{"testdata/69_core.flyte_basics.lp.greet_1.pb"}
+		var registerResults []Result
+		results, err := registerFile(ctx, args[0], registerResults, cmdCtx)
+		assert.Equal(t, 1, len(results))
+		assert.Equal(t, "Failed", results[0].Status)
+		assert.Equal(t, "Error registering file due to rpc error: code = InvalidArgument desc = Invalid", results[0].Info)
+		assert.NotNil(t, err)
+	})
+}
diff --git a/flytectl/cmd/register/testdata/69_core.flyte_basics.lp.greet_1.pb b/flytectl/cmd/register/testdata/69_core.flyte_basics.lp.greet_1.pb
diff --git a/flytectl/main.go b/flytectl/main.go
@@ -1,13 +1,16 @@
 package main
 
 import (
-	"fmt"
+	"context"
+	"os"
 
 	"github.com/flyteorg/flytectl/cmd"
+	"github.com/flyteorg/flytestdlib/logger"
 )
 
 func main() {
 	if err := cmd.ExecuteCmd(); err != nil {
-		fmt.Printf("error: %v\n", err)
+		logger.Error(context.TODO(), err)
+		os.Exit(1)
 	}
 }