workflow_test.go

package blackstart

import (
	"context"
	"errors"
	"reflect"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var cleanupModuleCloseCalls atomic.Int32

type cleanupTestModule struct{}

func init() {
	RegisterModule("cleanup_test_module", func() Module { return &cleanupTestModule{} })
}

func (m *cleanupTestModule) Info() ModuleInfo {
	return ModuleInfo{
		Id: "cleanup_test_module",
		Inputs: map[string]InputValue{
			testCheckResult: {
				Type:     reflect.TypeFor[bool](),
				Required: true,
			},
			testSetResult: {
				Type:     reflect.TypeFor[bool](),
				Required: true,
			},
		},
		Outputs: map[string]OutputValue{},
	}
}

func (m *cleanupTestModule) Validate(_ Operation) error { return nil }
func (m *cleanupTestModule) Check(ctx ModuleContext) (bool, error) {
	return InputAs[bool](ctxMustInput(ctx, testCheckResult), true)
}
func (m *cleanupTestModule) Set(_ ModuleContext) error { return nil }
func (m *cleanupTestModule) Close() error {
	cleanupModuleCloseCalls.Add(1)
	return nil
}

func ctxMustInput(ctx ModuleContext, key string) Input {
	in, _ := ctx.Input(key)
	return in
}

type testOpValues map[string]map[string]interface{}

func TestWorkflowExecution(t *testing.T) {
	tests := []struct {
		// name is the subtest name.
		name string

		// wf is the test Workflow to be executed. Currently, we only have the 'test_module' module
		// that is used in the test workflows.
		wf Workflow

		// checkOpInputs is a map of operation id to a map of input key to input value. This is used to
		// check the inputs to the specified operations in the test Workflow with expected values.
		checkOpInputs *testOpValues

		// checkOpOutputs is a map of operation id to a map of output key to output value. This is used to
		// check the outputs from the specified operations in the test Workflow with expected values.
		checkOpOutputs *testOpValues
	}{
		{
			name: "check_true_set_false",
			wf: Workflow{
				Name:        "",
				Description: "",
				Operations: []Operation{
					{
						Module:      "test_module",
						Id:          "test0",
						Name:        "test module",
						Description: "foo",
						DependsOn:   nil,
						Inputs: map[string]Input{
							testCheckResult: NewInputFromValue(true),
							testSetResult:   NewInputFromValue(false),
						},
						DoesNotExist: false,
						Tainted:      false,
					},
				},
			},
			checkOpInputs: &testOpValues{
				"test0": {
					testCheckResult: true,
				},
			},
			checkOpOutputs: &testOpValues{
				"test0": {
					testSetResult: false,
				},
			},
		},
		{
			name: "check_true_set_true",
			wf: Workflow{
				Name:        "",
				Description: "",
				Operations: []Operation{
					{
						Module:      "test_module",
						Id:          "test0",
						Name:        "test module",
						Description: "foo",
						DependsOn:   nil,
						Inputs: map[string]Input{
							testCheckResult: NewInputFromValue(true),
							testSetResult:   NewInputFromValue(true),
						},
						DoesNotExist: false,
						Tainted:      false,
					},
				},
			},
			checkOpInputs: &testOpValues{
				"test0": {
					testCheckResult: true,
				},
			},
			checkOpOutputs: &testOpValues{
				"test0": {
					testSetResult: true,
				},
			},
		},
		{
			name: "check_false_set_false",
			wf: Workflow{
				Name:        "",
				Description: "",
				Operations: []Operation{
					{
						Module:      "test_module",
						Id:          "test0",
						Name:        "test module",
						Description: "foo",
						DependsOn:   nil,
						Inputs: map[string]Input{
							testCheckResult: NewInputFromValue(false),
							testSetResult:   NewInputFromValue(false),
						},
						DoesNotExist: false,
						Tainted:      false,
					},
				},
			},
			checkOpInputs: &testOpValues{
				"test0": {
					testCheckResult: false,
				},
			},
			checkOpOutputs: &testOpValues{
				"test0": {
					testSetResult: false,
				},
			},
		},
		{
			name: "check_false_set_true",
			wf: Workflow{
				Name:        "",
				Description: "",
				Operations: []Operation{
					{
						Module:      "test_module",
						Id:          "test0",
						Name:        "test module",
						Description: "foo",
						DependsOn:   nil,
						Inputs: map[string]Input{
							testCheckResult: NewInputFromValue(false),
							testSetResult:   NewInputFromValue(true),
						},
						DoesNotExist: false,
						Tainted:      false,
					},
				},
			},
			checkOpInputs: &testOpValues{
				"test0": {
					testCheckResult: false,
				},
			},
			checkOpOutputs: &testOpValues{
				"test0": {
					testSetResult: true,
				},
			},
		},
		{
			name: "check_foo_set_bar",
			wf: Workflow{
				Name:        "",
				Description: "",
				Operations: []Operation{
					{
						Module:      "test_module",
						Id:          "test0",
						Name:        "test module",
						Description: "foo",
						DependsOn:   nil,
						Inputs: map[string]Input{
							"input_0":       NewInputFromValue("foo"),
							testCheckResult: NewInputFromValue(true),
							testSetResult:   NewInputFromValue(true),
						},
						DoesNotExist: false,
						Tainted:      false,
					},
				},
			},
			checkOpInputs: &testOpValues{
				"test0": {
					"input_0": "foo",
				},
			},
			checkOpOutputs: &testOpValues{
				"test0": {
					"output_0": "bar",
				},
			},
		},
	}

	for _, tt := range tests {
		t.Run(
			tt.name, func(t *testing.T) {
				ctx := context.Background()
				logger := NewTestingLogger()
				we := newWorkflowExecution(&tt.wf, logger)
				res := we.execute(ctx)
				assert.NoError(t, res.Err)

				// check inputs to specified operations in the test Workflow
				if tt.checkOpInputs != nil {
					for opId, iv := range *tt.checkOpInputs {
						for k, v := range iv {
							var input Input
							var err error
							input, err = we.opCtxs[opId].Input(k)
							assert.NoError(t, err)

							// normally, a module would need to implement the value conversion for inputs.
							switch v.(type) {
							case bool:
								convertedInput, convErr := InputAs[bool](input, true)
								assert.NoError(t, convErr)
								assert.Equal(t, v, convertedInput)
							case string:
								convertedInput, convErr := InputAs[string](input, true)
								assert.NoError(t, convErr)
								assert.Equal(t, v, convertedInput)
							}
						}
					}
				}

				// check outputs from specified operations in the test Workflow
				if tt.checkOpOutputs != nil {
					for opId, iv := range *tt.checkOpOutputs {
						for k, v := range iv {
							var output interface{}
							output, _ = we.opCtxs[opId].getOutput(k)

							// normally, a module would need to implement the value conversion for inputs, here
							// we need to do it for the outputs just for test convenience.
							var ok bool
							switch v.(type) {
							case bool:
								var convertedOutput bool
								convertedOutput, ok = output.(bool)
								assert.True(t, ok)
								assert.Equal(t, v, convertedOutput)
							case string:
								var convertedOutput string
								convertedOutput, ok = output.(string)
								assert.True(t, ok)
								assert.Equal(t, v, convertedOutput)
							}

						}
					}
				}
			},
		)
	}
}

func TestWorkflowExecution_DuplicateOperationID(t *testing.T) {
	wf := Workflow{
		Name: "duplicate-op-id",
		Operations: []Operation{
			{
				Id:     "dup",
				Module: "test_module",
				Inputs: map[string]Input{
					testCheckResult: NewInputFromValue(true),
					testSetResult:   NewInputFromValue(true),
				},
			},
			{
				Id:     "dup",
				Module: "test_module",
				Inputs: map[string]Input{
					testCheckResult: NewInputFromValue(true),
					testSetResult:   NewInputFromValue(true),
				},
			},
		},
	}

	res := wf.Run(context.Background())
	require.Error(t, res.Err)
	require.Contains(t, res.Err.Error(), `duplicate operation id "dup"`)
	require.NotNil(t, res.Op)
	require.Equal(t, "dup", res.Op.Id)
	require.Equal(t, phaseSetup, res.Phase)
}

func TestWorkflowExecution_CallsOptionalModuleClose(t *testing.T) {
	cleanupModuleCloseCalls.Store(0)
	wf := Workflow{
		Name: "cleanup-close-test",
		Operations: []Operation{
			{
				Id:     "cleanup-op",
				Module: "cleanup_test_module",
				Inputs: map[string]Input{
					testCheckResult: NewInputFromValue(true),
					testSetResult:   NewInputFromValue(true),
				},
			},
		},
	}

	res := wf.Run(context.Background())
	require.NoError(t, res.Err)
	require.Equal(t, int32(1), cleanupModuleCloseCalls.Load())
}

// TestOpoSort tests the topological sorting of operations into an expected order.
func TestOpoSort(t *testing.T) {
	tests := []struct {
		name   string
		in     []Operation
		out    []string
		outErr error
	}{
		{
			name: "sort_0",
			in: []Operation{
				{
					Id:        "test0",
					DependsOn: nil,
				},
				{
					Id:        "test1",
					DependsOn: []string{"test0"},
				},
				{
					Id:        "test2",
					DependsOn: []string{"test1"},
				},
			},
			out: []string{"test0", "test1", "test2"},
		},
		{
			name: "sort_1",
			in: []Operation{
				{
					Id:        "test0",
					DependsOn: nil,
				},
				{
					Id:        "test1",
					DependsOn: []string{"test2"},
				},
				{
					Id:        "test2",
					DependsOn: []string{"test0"},
				},
			},
			out: []string{"test0", "test2", "test1"},
		},
		// multi-level
		{
			name: "sort_2",
			in: []Operation{
				{
					Id:        "test0",
					DependsOn: nil,
				},
				{
					Id:        "test1",
					DependsOn: []string{"test0"},
				},
				{
					Id:        "test2",
					DependsOn: []string{"test0"},
				},
				{
					Id:        "test3",
					DependsOn: []string{"test2"},
				},
			},
			out: []string{"test0", "test1", "test2", "test3"},
		},
		// Self-cycle
		{
			name: "sort_3",
			in: []Operation{
				{
					Id:        "test0",
					DependsOn: nil,
				},
				{
					Id:        "test1",
					DependsOn: []string{"test0"},
				},
				{
					Id:        "test2",
					DependsOn: []string{"test0"},
				},
				{
					Id:        "test3",
					DependsOn: []string{"test3"},
				},
			},
			out:    nil,
			outErr: ErrOperationCycle,
		},
		// Circle
		{
			name: "sort_4",
			in: []Operation{
				{
					Id:        "test0",
					DependsOn: []string{"test3"},
				},
				{
					Id:        "test1",
					DependsOn: []string{"test0"},
				},
				{
					Id:        "test2",
					DependsOn: []string{"test1"},
				},
				{
					Id:        "test3",
					DependsOn: []string{"test2"},
				},
			},
			out:    nil,
			outErr: ErrOperationCycle,
		},
		// internal cycle
		{
			name: "sort_5",
			in: []Operation{
				{
					Id:        "test0",
					DependsOn: nil,
				},
				{
					Id:        "test1",
					DependsOn: []string{"test0", "test3"},
				},
				{
					Id:        "test2",
					DependsOn: []string{"test1"},
				},
				{
					Id:        "test3",
					DependsOn: []string{"test2"},
				},
			},
			out:    nil,
			outErr: ErrOperationCycle,
		},
		// multiple start nodes
		{
			name: "sort_6",
			in: []Operation{
				{
					Id:        "test0",
					DependsOn: nil,
				},
				{
					Id:        "test1",
					DependsOn: []string{"test0", "test3"},
				},
				{
					Id:        "test2",
					DependsOn: nil,
				},
				{
					Id:        "test3",
					DependsOn: []string{"test2"},
				},
			},
			out:    []string{"test0", "test2", "test3", "test1"},
			outErr: nil,
		},
	}

	for _, tt := range tests {
		t.Run(
			tt.name, func(t *testing.T) {
				result, err := opoSort(tt.in)
				assert.Equal(t, tt.out, result)
				if tt.outErr != nil {
					assert.True(t, errors.Is(err, tt.outErr))
				}
			},
		)
	}
}

func TestCheckInputsOutputs_StaticInputAcceptsAnyDeclaredType(t *testing.T) {
	op := &Operation{
		Id: "test-op",
		Inputs: map[string]Input{
			"value": NewInputFromValue(true),
		},
	}
	info := ModuleInfo{
		Inputs: map[string]InputValue{
			"value": {
				Required: true,
				Types:    []reflect.Type{reflect.TypeFor[string](), reflect.TypeFor[bool]()},
			},
		},
	}

	err := checkInputsOutputs(op, info, nil)
	require.NoError(t, err)
}

func TestCheckInputsOutputs_SupportsLegacySingleTypeField(t *testing.T) {
	op := &Operation{
		Id: "test-op",
		Inputs: map[string]Input{
			"value": NewInputFromValue("ok"),
		},
	}
	info := ModuleInfo{
		Inputs: map[string]InputValue{
			"value": {
				Required: true,
				Type:     reflect.TypeFor[string](),
			},
		},
	}

	err := checkInputsOutputs(op, info, nil)
	require.NoError(t, err)
}

func TestCheckInputsOutputs_DependencyOutputMatchesAnyDeclaredType(t *testing.T) {
	op := &Operation{
		Id: "test-op",
		Inputs: map[string]Input{
			"value": NewInputFromDep("dep-op", "result"),
		},
	}
	info := ModuleInfo{
		Inputs: map[string]InputValue{
			"value": {
				Required: true,
				Types:    []reflect.Type{reflect.TypeFor[string](), reflect.TypeFor[bool]()},
			},
		},
	}
	opsInfo := map[string]ModuleInfo{
		"dep-op": {
			Outputs: map[string]OutputValue{
				"result": {Type: reflect.TypeFor[bool]()},
			},
		},
	}

	err := checkInputsOutputs(op, info, opsInfo)
	require.NoError(t, err)
}

func TestCheckInputsOutputs_DependencyOutputTypeMismatch(t *testing.T) {
	op := &Operation{
		Id: "test-op",
		Inputs: map[string]Input{
			"value": NewInputFromDep("dep-op", "result"),
		},
	}
	info := ModuleInfo{
		Inputs: map[string]InputValue{
			"value": {
				Required: true,
				Types:    []reflect.Type{reflect.TypeFor[string](), reflect.TypeFor[bool]()},
			},
		},
	}
	opsInfo := map[string]ModuleInfo{
		"dep-op": {
			Outputs: map[string]OutputValue{
				"result": {Type: reflect.TypeFor[int]()},
			},
		},
	}

	err := checkInputsOutputs(op, info, opsInfo)
	require.Error(t, err)
	require.ErrorContains(t, err, "expected type(s)")
}

func TestCheckInputsOutputs_StaticAnySliceMatchesStringSlice(t *testing.T) {
	op := &Operation{
		Id: "test-op",
		Inputs: map[string]Input{
			"value": NewInputFromValue([]any{"SELECT", "UPDATE"}),
		},
	}
	info := ModuleInfo{
		Inputs: map[string]InputValue{
			"value": {
				Required: true,
				Type:     reflect.TypeFor[[]string](),
			},
		},
	}

	err := checkInputsOutputs(op, info, nil)
	require.NoError(t, err)
}

func TestCheckInputsOutputs_StaticAnySliceMismatchStringSlice(t *testing.T) {
	op := &Operation{
		Id: "test-op",
		Inputs: map[string]Input{
			"value": NewInputFromValue([]any{"SELECT", 1}),
		},
	}
	info := ModuleInfo{
		Inputs: map[string]InputValue{
			"value": {
				Required: true,
				Type:     reflect.TypeFor[[]string](),
			},
		},
	}

	err := checkInputsOutputs(op, info, nil)
	require.Error(t, err)
	require.ErrorContains(t, err, "not assignable")
}