Fix spark tests

Signed-off-by: Fabio Graetz <[email protected]>

flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go

@@ -3,6 +3,7 @@ package spark
 import (
 	"context"
 	"os"
+	"reflect"
 	"strconv"
 	"testing"
 
@@ -118,9 +119,14 @@ func TestGetEventInfo(t *testing.T) {
 	assert.Equal(t, expectedLinks, generatedLinks)
 
 	info, err = getEventInfoForSpark(taskCtx, dummySparkApplication(sj.SubmittedState))
+	generatedLinks = make([]string, 0, len(info.Logs))
+	for _, l := range info.Logs {
+		generatedLinks = append(generatedLinks, l.Uri)
+	}
 	assert.NoError(t, err)
-	assert.Len(t, info.Logs, 1)
-	assert.Equal(t, "https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logStream:group=/kubernetes/flyte;prefix=var.log.containers.spark-app-name;streamFilter=typeLogStreamPrefix", info.Logs[0].Uri)
+	assert.Len(t, info.Logs, 5)
+	assert.Equal(t, expectedLinks[:5], generatedLinks) // No Spark Driver UI for Submitted state
+	assert.True(t, info.Logs[4].ShowWhilePending)      // All User Logs should be shown while pending
 
 	assert.NoError(t, setSparkConfig(&Config{
 		SparkHistoryServerURL: "spark-history.flyte",
@@ -406,6 +412,19 @@ func dummySparkTaskContext(taskTemplate *core.TaskTemplate, interruptible bool)
 	taskExecutionMetadata.On("GetOverrides").Return(overrides)
 	taskExecutionMetadata.On("GetK8sServiceAccount").Return("new-val")
 	taskCtx.On("TaskExecutionMetadata").Return(taskExecutionMetadata)
+
+	inputState := k8s.PluginState{}
+	pluginStateReaderMock := mocks.PluginStateReader{}
+	pluginStateReaderMock.On("Get", mock.AnythingOfType(reflect.TypeOf(&inputState).String())).Return(
+		func(v interface{}) uint8 {
+			*(v.(*k8s.PluginState)) = inputState
+			return 0
+		},
+		func(v interface{}) error {
+			return nil
+		})
+
+	taskCtx.OnPluginStateReader().Return(&pluginStateReaderMock)
 	return taskCtx
 }