Skip to content

Commit

Permalink
Adopt flyteidl's ordered variable map change (flyteorg#198)
Browse files Browse the repository at this point in the history
  • Loading branch information
mayitbeegh authored Sep 2, 2021
1 parent 2e8a22b commit d76eb15
Show file tree
Hide file tree
Showing 18 changed files with 340 additions and 111 deletions.
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ require (
github.com/aws/aws-sdk-go-v2/config v1.0.0
github.com/aws/aws-sdk-go-v2/service/athena v1.0.0
github.com/coocood/freecache v1.1.1
github.com/flyteorg/flyteidl v0.19.2
github.com/flyteorg/flyteidl v0.20.0
github.com/flyteorg/flytestdlib v0.3.33
github.com/go-logr/zapr v0.4.0 // indirect
github.com/go-test/deep v1.0.7
Expand Down
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -226,8 +226,8 @@ github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5Kwzbycv
github.com/fatih/color v1.10.0 h1:s36xzo75JdqLaaWoiEHk767eHiwo0598uUxyfiPkDsg=
github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGEBuJM=
github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
github.com/flyteorg/flyteidl v0.19.2 h1:jXuRrLJEzSo33N9pw7bMEd6mRYSL7LCz/vnazz5XcOg=
github.com/flyteorg/flyteidl v0.19.2/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
github.com/flyteorg/flyteidl v0.20.0 h1:g5xGayFfPSzFJxJedgL390WFSEbGYjFiPey+NXAB030=
github.com/flyteorg/flyteidl v0.20.0/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
github.com/flyteorg/flytestdlib v0.3.13/go.mod h1:Tz8JCECAbX6VWGwFT6cmEQ+RJpZ/6L9pswu3fzWs220=
github.com/flyteorg/flytestdlib v0.3.33 h1:+oCx3zXUIldL7CWmNMD7PMFPXvGqaPgYkSKn9wB6qvY=
github.com/flyteorg/flytestdlib v0.3.33/go.mod h1:7cDWkY3v7xsoesFcDdu6DSW5Q2U2W5KlHUbUHSwBG1Q=
Expand Down
182 changes: 141 additions & 41 deletions go/tasks/pluginmachinery/flytek8s/copilot_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,19 @@ func TestDownloadCommandArgs(t *testing.T) {
assert.Error(t, err)

iFace := &core.VariableMap{
Variables: map[string]*core.Variable{
"x": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
"y": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "x",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
{
Name: "y",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
}
d, err := DownloadCommandArgs("s3://from", "s3://output-meta", "/to", core.DataLoadingConfig_JSON, iFace)
Expand All @@ -121,10 +131,10 @@ func TestDownloadCommandArgs(t *testing.T) {
vm := &core.VariableMap{}
assert.NoError(t, proto.Unmarshal(serIFaceBytes, vm))
assert.Len(t, vm.Variables, 2)
for k, v := range iFace.Variables {
v2, ok := vm.Variables[k]
assert.True(t, ok)
assert.Equal(t, v.Type.GetSimple(), v2.Type.GetSimple(), "for %s, types do not match", k)
for i, v := range iFace.Variables {
v2 := vm.Variables[i]
assert.Equal(t, v.Name, v2.Name, "for index %d, keys do not match", i)
assert.Equal(t, v.Var.Type.GetSimple(), v2.Var.Type.GetSimple(), "for %s, types do not match", v.Name)
}
}
}
Expand All @@ -136,9 +146,19 @@ func TestSidecarCommandArgs(t *testing.T) {

iFace := &core.TypedInterface{
Outputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"x": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
"y": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "x",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
{
Name: "y",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand All @@ -156,10 +176,10 @@ func TestSidecarCommandArgs(t *testing.T) {
if2 := &core.TypedInterface{}
assert.NoError(t, proto.Unmarshal(serIFaceBytes, if2))
assert.Len(t, if2.Outputs.Variables, 2)
for k, v := range iFace.Outputs.Variables {
v2, ok := if2.Outputs.Variables[k]
assert.True(t, ok)
assert.Equal(t, v.Type.GetSimple(), v2.Type.GetSimple(), "for %s, types do not match", k)
for i, v := range iFace.Outputs.Variables {
v2 := if2.Outputs.Variables[i]
assert.Equal(t, v.Name, v2.Name, "for index %d, keys do not match", i)
assert.Equal(t, v.Var.Type.GetSimple(), v2.Var.Type.GetSimple(), "for %s, types do not match", v.Name)
}
}
}
Expand Down Expand Up @@ -358,14 +378,29 @@ func TestAddCoPilotToContainer(t *testing.T) {
c := v1.Container{}
iface := &core.TypedInterface{
Inputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"x": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
"y": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "x",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
{
Name: "y",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
Outputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"o": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "o",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand All @@ -380,14 +415,29 @@ func TestAddCoPilotToContainer(t *testing.T) {
c := v1.Container{}
iface := &core.TypedInterface{
Inputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"x": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
"y": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "x",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
{
Name: "y",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
Outputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"o": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "o",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand All @@ -406,9 +456,19 @@ func TestAddCoPilotToContainer(t *testing.T) {
c := v1.Container{}
iface := &core.TypedInterface{
Inputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"x": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
"y": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "x",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
{
Name: "y",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand All @@ -427,8 +487,13 @@ func TestAddCoPilotToContainer(t *testing.T) {
c := v1.Container{}
iface := &core.TypedInterface{
Outputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"o": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "o",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand Down Expand Up @@ -507,14 +572,29 @@ func TestAddCoPilotToPod(t *testing.T) {
pod := v1.PodSpec{}
iface := &core.TypedInterface{
Inputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"x": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
"y": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "x",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
{
Name: "y",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
Outputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"o": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "o",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand Down Expand Up @@ -544,9 +624,19 @@ func TestAddCoPilotToPod(t *testing.T) {
pod := v1.PodSpec{}
iface := &core.TypedInterface{
Inputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"x": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
"y": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "x",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
{
Name: "y",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand All @@ -564,8 +654,13 @@ func TestAddCoPilotToPod(t *testing.T) {
pod := v1.PodSpec{}
iface := &core.TypedInterface{
Outputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"o": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "o",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand All @@ -583,8 +678,13 @@ func TestAddCoPilotToPod(t *testing.T) {
pod := v1.PodSpec{}
iface := &core.TypedInterface{
Outputs: &core.VariableMap{
Variables: map[string]*core.Variable{
"o": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}},
Variables: []*core.VariableMapEntry{
{
Name: "o",
Var: &core.Variable{
Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}},
},
},
},
},
}
Expand Down
25 changes: 25 additions & 0 deletions go/tasks/pluginmachinery/utils/task.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
package utils

import "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"

// This function extracts the variable named "results" from the TaskTemplate
func GetResultsVariable(taskTemplate *core.TaskTemplate) (results *core.Variable, exists bool) {
if taskTemplate == nil {
return nil, false
}
if taskTemplate.Interface == nil {
return nil, false
}
if taskTemplate.Interface.Outputs == nil {
return nil, false
}
if taskTemplate.Interface.Outputs.Variables == nil {
return nil, false
}
for _, e := range taskTemplate.Interface.Outputs.Variables {
if e.Name == "results" {
return e.Var, true
}
}
return nil, false
}
54 changes: 54 additions & 0 deletions go/tasks/pluginmachinery/utils/task_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package utils

import (
"testing"

"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
"github.com/stretchr/testify/assert"
)

func TestGetResultsVariable(t *testing.T) {
t.Run("variable not found with nil pointer", func(t *testing.T) {
emptyTaskTemplate := &core.TaskTemplate{
Interface: nil,
}
_, found := GetResultsVariable(emptyTaskTemplate)
assert.False(t, found)
})

t.Run("variable not found", func(t *testing.T) {
taskTemplate := &core.TaskTemplate{
Interface: &core.TypedInterface{
Outputs: &core.VariableMap{
Variables: []*core.VariableMapEntry{
{
Name: "o0",
},
},
},
},
}
_, found := GetResultsVariable(taskTemplate)
assert.False(t, found)
})

t.Run("happy case", func(t *testing.T) {
taskTemplate := &core.TaskTemplate{
Interface: &core.TypedInterface{
Outputs: &core.VariableMap{
Variables: []*core.VariableMapEntry{
{
Name: "results",
Var: &core.Variable{
Description: "athena result",
},
},
},
},
},
}
v, found := GetResultsVariable(taskTemplate)
assert.True(t, found)
assert.Equal(t, "athena result", v.Description)
})
}
Loading

0 comments on commit d76eb15

Please sign in to comment.