Add support for gate nodes (#399)

* Add support for gate nodes Signed-off-by: eduardo apolinario <[email protected]> * Linting Signed-off-by: eduardo apolinario <[email protected]> --------- Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: eduardo apolinario <[email protected]>
flyteorg · Mar 14, 2023 · 29da288 · 29da288
1 parent e55c649
commit 29da288
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/cmd/register/register_util.go b/cmd/register/register_util.go
@@ -245,6 +245,8 @@ func hydrateNode(node *core.Node, version string, force bool) error {
 		default:
 			return fmt.Errorf("unknown type %T", branchNodeWrapper.BranchNode.IfElse.Default)
 		}
+	case *core.Node_GateNode:
+		// Do nothing.
 	default:
 		return fmt.Errorf("unknown type %T", v)
 	}

diff --git a/cmd/register/register_util_test.go b/cmd/register/register_util_test.go
@@ -35,6 +35,7 @@ import (
 	"github.com/stretchr/testify/mock"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/durationpb"
 )
 
 type MockHTTPClient struct {
@@ -501,6 +502,64 @@ func TestHydrateNode(t *testing.T) {
 	})
 }
 
+func TestHydrateGateNode(t *testing.T) {
+	t.Run("Hydrate Sleep", func(t *testing.T) {
+		registerFilesSetup()
+		// Write a node that contains a GateNode
+		node := &core.Node{
+			Target: &core.Node_GateNode{
+				GateNode: &core.GateNode{
+					Condition: &core.GateNode_Sleep{
+						Sleep: &core.SleepCondition{
+							Duration: &durationpb.Duration{
+								Seconds: 10,
+							},
+						},
+					},
+				},
+			},
+		}
+		err := hydrateNode(node, rconfig.DefaultFilesConfig.Version, true)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Hydrate Signal", func(t *testing.T) {
+		registerFilesSetup()
+		// Write a node that contains a GateNode
+		node := &core.Node{
+			Target: &core.Node_GateNode{
+				GateNode: &core.GateNode{
+					Condition: &core.GateNode_Signal{
+						Signal: &core.SignalCondition{
+							SignalId: "abc",
+						},
+					},
+				},
+			},
+		}
+		err := hydrateNode(node, rconfig.DefaultFilesConfig.Version, true)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Hydrate Approve", func(t *testing.T) {
+		registerFilesSetup()
+		// Write a node that contains a GateNode
+		node := &core.Node{
+			Target: &core.Node_GateNode{
+				GateNode: &core.GateNode{
+					Condition: &core.GateNode_Approve{
+						Approve: &core.ApproveCondition{
+							SignalId: "abc",
+						},
+					},
+				},
+			},
+		}
+		err := hydrateNode(node, rconfig.DefaultFilesConfig.Version, true)
+		assert.Nil(t, err)
+	})
+}
+
 func TestHydrateTaskSpec(t *testing.T) {
 	testScope := promutils.NewTestScope()
 	labeled.SetMetricKeys(contextutils.AppNameKey, contextutils.ProjectKey, contextutils.DomainKey)