Add unit test

flyteorg · Aug 19, 2024 · 503cbd3 · 503cbd3
1 parent 252bac0
commit 503cbd3
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/plugin_test.go b/flyteplugins/go/tasks/plugins/webapi/agent/plugin_test.go
@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"golang.org/x/exp/maps"
+	"google.golang.org/protobuf/types/known/structpb"
 
 	agentMocks "github.com/flyteorg/flyte/flyteidl/clients/go/admin/mocks"
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
@@ -114,17 +115,24 @@ func TestPlugin(t *testing.T) {
 	})
 
 	t.Run("test RUNNING Status", func(t *testing.T) {
+		simpleStruct := structpb.Struct{
+			Fields: map[string]*structpb.Value{
+				"foo": {Kind: &structpb.Value_StringValue{StringValue: "foo"}},
+			},
+		}
 		taskContext := new(webapiPlugin.StatusContext)
 		taskContext.On("Resource").Return(ResourceWrapper{
-			State:    admin.State_RUNNING,
-			Outputs:  nil,
-			Message:  "Job is running",
-			LogLinks: []*flyteIdlCore.TaskLog{{Uri: "http://localhost:3000/log", Name: "Log Link"}},
+			State:      admin.State_RUNNING,
+			Outputs:    nil,
+			Message:    "Job is running",
+			LogLinks:   []*flyteIdlCore.TaskLog{{Uri: "http://localhost:3000/log", Name: "Log Link"}},
+			CustomInfo: &simpleStruct,
 		})
 
 		phase, err := plugin.Status(context.Background(), taskContext)
 		assert.NoError(t, err)
 		assert.Equal(t, pluginsCore.PhaseRunning, phase.Phase())
+		assert.Equal(t, &simpleStruct, phase.Info().CustomInfo)
 	})
 
 	t.Run("test PERMANENT_FAILURE Status", func(t *testing.T) {