tasks/chasm/field_test.go

package chasm

import (
	"context"
	"reflect"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	persistencespb "github.com/hanzoai/tasks/api/persistence/v1"
	"github.com/hanzoai/tasks/common/clock"
	"github.com/hanzoai/tasks/common/definition"
	"github.com/hanzoai/tasks/common/log"
	"github.com/hanzoai/tasks/common/metrics"
	"github.com/hanzoai/tasks/common/primitives"
	"github.com/hanzoai/tasks/common/testing/protorequire"
	"github.com/hanzoai/tasks/common/testing/testlogger"
	"go.uber.org/mock/gomock"
)

type fieldSuite struct {
	suite.Suite
	*require.Assertions
	protorequire.ProtoAssertions

	controller  *gomock.Controller
	nodeBackend *MockNodeBackend

	registry        *Registry
	timeSource      *clock.EventTimeSource
	nodePathEncoder NodePathEncoder
	logger          log.Logger
	metricsHandler  metrics.Handler
}

func TestFieldSuite(t *testing.T) {
	suite.Run(t, new(fieldSuite))
}

func (s *fieldSuite) SetupTest() {
	s.initAssertions()
	s.controller = gomock.NewController(s.T())
	s.nodeBackend = &MockNodeBackend{}

	s.logger = testlogger.NewTestLogger(s.T(), testlogger.FailOnAnyUnexpectedError)
	s.metricsHandler = metrics.NoopMetricsHandler
	s.registry = NewRegistry(s.logger)
	err := s.registry.Register(newTestLibrary(s.controller))
	s.NoError(err)

	s.timeSource = clock.NewEventTimeSource()
	s.nodePathEncoder = &testNodePathEncoder{}
}

func (s *fieldSuite) SetupSubTest() {
	s.initAssertions()
}

func (s *fieldSuite) initAssertions() {
	// `s.Assertions` (as well as other test helpers which depends on `s.T()`) must be initialized on
	// both test and subtest levels (but not suite level, where `s.T()` is `nil`).
	//
	// If these helpers are not reinitialized on subtest level, any failed `assert` in
	// subtest will fail the entire test (not subtest) immediately without running other subtests.

	s.Assertions = require.New(s.T())
	s.ProtoAssertions = protorequire.New(s.T())
}

func (s *fieldSuite) TestInternalFieldName() {
	f := Field[any]{}
	fT := reflect.TypeOf(f)

	_, ok := fT.FieldByName(internalFieldName)
	s.True(ok, "expected field %s not found", internalFieldName)
}

func (s *fieldSuite) TestFieldGetSimple() {
	tests := []struct {
		name     string
		field    Field[*TestSubComponent1]
		expected *TestSubComponent1
	}{
		{
			name: "Get with non-nil value",
			field: Field[*TestSubComponent1]{
				Internal: newFieldInternalWithValue(
					fieldTypeComponent,
					&TestSubComponent1{SubComponent1Data: &protoMessageType{
						CreateRequestId: "component-data",
					}},
				)},
			expected: &TestSubComponent1{SubComponent1Data: &protoMessageType{
				CreateRequestId: "component-data",
			}},
		},
		{
			name: "Get with nil value and nil node",
			field: Field[*TestSubComponent1]{
				Internal: newFieldInternalWithNode(nil),
			},
			expected: nil,
		},
	}

	for _, tt := range tests {
		s.Run(tt.name, func() {
			result, _ := tt.field.TryGet(nil)
			s.Equal(tt.expected, result)
		})
	}
}

func (s *fieldSuite) TestFieldGetComponent() {
	serializedNodes := testComponentSerializedNodes()

	node, err := s.newTestTree(serializedNodes)
	s.NoError(err)

	chasmContext := NewMutableContext(context.Background(), node)

	c, err := node.Component(chasmContext, ComponentRef{componentPath: rootPath})
	s.NoError(err)
	s.NotNil(c)

	tc := c.(*TestComponent)

	sc1 := tc.SubComponent1.Get(chasmContext)
	s.NotNil(sc1)
	s.ProtoEqual(&protoMessageType{
		CreateRequestId: "sub-component1-data",
	}, sc1.SubComponent1Data)

	sd1 := tc.SubData1.Get(chasmContext)
	s.NotNil(sd1)
	s.ProtoEqual(&protoMessageType{
		CreateRequestId: "sub-data1",
	}, sd1)
}

func (s *fieldSuite) newTestTree(
	serializedNodes map[string]*persistencespb.ChasmNode,
) (*Node, error) {
	if len(serializedNodes) == 0 {
		return NewEmptyTree(
			s.registry,
			s.timeSource,
			s.nodeBackend,
			s.nodePathEncoder,
			s.logger,
			s.metricsHandler,
		), nil
	}
	return NewTreeFromDB(
		serializedNodes,
		s.registry,
		s.timeSource,
		s.nodeBackend,
		s.nodePathEncoder,
		s.logger,
		s.metricsHandler,
	)
}

// setupComponentWithTree creates a basic component structure and attaches it to the tree.
func (s *fieldSuite) setupComponentWithTree(rootComponent *TestComponent) (*Node, MutableContext, error) {
	rootNode := NewEmptyTree(
		s.registry,
		s.timeSource,
		s.nodeBackend,
		s.nodePathEncoder,
		s.logger,
		s.metricsHandler,
	)
	if err := rootNode.SetRootComponent(rootComponent); err != nil {
		return nil, nil, err
	}

	return rootNode, NewMutableContext(context.Background(), rootNode), nil
}

func (s *fieldSuite) TestDeferredPointerResolution() {
	workflowKey := definition.NewWorkflowKey(
		primitives.NewUUID().String(),
		primitives.NewUUID().String(),
		primitives.NewUUID().String(),
	)
	s.nodeBackend = &MockNodeBackend{
		HandleNextTransitionCount: func() int64 { return 1 },
		HandleGetCurrentVersion:   func() int64 { return 1 },
		HandleGetWorkflowKey:      func() definition.WorkflowKey { return workflowKey },
	}

	// Create component structure that will simulate StartExecution scenario.
	sc2 := &TestSubComponent2{
		SubComponent2Data: &protoMessageType{
			CreateRequestId: "sub-component2-data",
		},
	}

	sc1 := &TestSubComponent1{
		SubComponent1Data: &protoMessageType{
			CreateRequestId: "sub-component1-data",
		},
	}

	rootComponent := &TestComponent{
		ComponentData: &protoMessageType{
			CreateRequestId: "component-data",
		},
		SubComponent1: NewComponentField(nil, sc1),
	}

	rootNode, ctx, err := s.setupComponentWithTree(rootComponent)
	s.NoError(err)

	// Get components from tree to mark nodes as needing sync.

	rootComponentInterface, err := rootNode.Component(ctx, ComponentRef{})
	s.NoError(err)
	rootComponent = rootComponentInterface.(*TestComponent)
	sc1 = rootComponent.SubComponent1.Get(ctx)

	// Create deferred pointers.
	sc1.SubComponent2Pointer = ComponentPointerTo(ctx, sc2)
	rootComponent.SubComponent2 = NewComponentField(nil, sc2)

	data := &protoMessageType{CreateRequestId: "sub-data-1"}
	sc1.DataPointer = DataPointerTo(ctx, data)
	rootComponent.SubData1 = NewDataField(ctx, data)

	// Verify it's a deferred pointer storing the component directly.
	s.Equal(fieldTypeDeferredPointer, sc1.SubComponent2Pointer.Internal.fieldType())
	s.Equal(fieldTypeDeferredPointer, sc1.DataPointer.Internal.fieldType())
	s.Equal(sc2, sc1.SubComponent2Pointer.Internal.v)
	s.Equal(data, sc1.DataPointer.Internal.v)

	// CloseTransaction should resolve the deferred pointer.
	mutations, err := rootNode.CloseTransaction()
	s.NoError(err)
	s.NotEmpty(mutations.UpdatedNodes)

	// Verify the pointers were resolved to a regular pointer with path.
	s.Equal(fieldTypePointer, sc1.SubComponent2Pointer.Internal.fieldType())
	s.Equal(fieldTypePointer, sc1.DataPointer.Internal.fieldType())

	cResolvedPath, ok := sc1.SubComponent2Pointer.Internal.v.([]string)
	s.True(ok)
	s.Equal([]string{"SubComponent2"}, cResolvedPath)
	dResolvedPath, ok := sc1.DataPointer.Internal.v.([]string)
	s.True(ok)
	s.Equal([]string{"SubData1"}, dResolvedPath)

	// Verify we can dereference the pointers.
	resolvedComponent := sc1.SubComponent2Pointer.Get(ctx)
	s.Equal(sc2, resolvedComponent)

	// TODO - this doesn't resolve, but I've manually verified the tree structure looks correct
	// TODO
	// TODO
	// resolvedData, err := sc1.DataPointer.Get(ctx)
	// s.NoError(err)
	// s.Equal(sc2.SubComponent2Data, resolvedData)
}

func (s *fieldSuite) TestMixedPointerScenario() {
	workflowKey := definition.NewWorkflowKey(
		primitives.NewUUID().String(),
		primitives.NewUUID().String(),
		primitives.NewUUID().String(),
	)
	s.nodeBackend = &MockNodeBackend{
		HandleNextTransitionCount: func() int64 { return 1 },
		HandleGetCurrentVersion:   func() int64 { return 1 },
		HandleGetWorkflowKey:      func() definition.WorkflowKey { return workflowKey },
	}

	existingComponent := &TestSubComponent11{
		SubComponent11Data: &protoMessageType{CreateRequestId: "existing-component"},
	}

	sc1 := &TestSubComponent1{
		SubComponent1Data: &protoMessageType{CreateRequestId: "sub-component1-data"},
		SubComponent11:    NewComponentField(nil, existingComponent),
	}

	rootComponent := &TestComponent{
		ComponentData: &protoMessageType{CreateRequestId: "component-data"},
		SubComponent1: NewComponentField(nil, sc1),
	}

	rootNode, ctx, err := s.setupComponentWithTree(rootComponent)
	s.NoError(err)

	// Get components from tree to mark nodes as needing sync.
	rootComponentInterface, err := rootNode.Component(ctx, ComponentRef{})
	s.NoError(err)
	rootComponent = rootComponentInterface.(*TestComponent)

	rootComponent.SubComponent11Pointer = ComponentPointerTo(ctx, existingComponent)

	// Close the transaction to resolve SubComponent11Pointer's field to existingComponent.
	_, err = rootNode.CloseTransaction()
	s.NoError(err)
	s.Equal(fieldTypePointer, rootComponent.SubComponent11Pointer.Internal.fieldType())

	// For a new transaction, get the components from the tree again,
	// otherwise those nodes will not be marked as dirty.

	ctx2 := NewMutableContext(context.Background(), rootNode)
	rootComponentInterface, err = rootNode.Component(ctx2, ComponentRef{})
	s.NoError(err)

	rootComponent = rootComponentInterface.(*TestComponent)
	sc1 = rootComponent.SubComponent1.Get(ctx2)

	// Now, add a new component and deferred pointer for it.
	newComponent := &TestSubComponent2{
		SubComponent2Data: &protoMessageType{CreateRequestId: "new-component"},
	}

	sc1.SubComponent2Pointer = ComponentPointerTo(ctx2, newComponent)

	// Now add the component to the tree so it can be resolved during CloseTransaction.
	rootComponent.SubComponent2 = NewComponentField(ctx, newComponent)

	s.Equal(fieldTypePointer, rootComponent.SubComponent11Pointer.Internal.fieldType())
	s.Equal(fieldTypeDeferredPointer, sc1.SubComponent2Pointer.Internal.fieldType())

	_, err = rootNode.CloseTransaction()
	s.NoError(err)

	// Ensure both pointers have been resolved.
	s.Equal(fieldTypePointer, rootComponent.SubComponent11Pointer.Internal.fieldType())
	s.Equal(fieldTypePointer, sc1.SubComponent2Pointer.Internal.fieldType())

	resolved1 := rootComponent.SubComponent11Pointer.Get(ctx2)
	s.Equal(existingComponent, resolved1)

	resolved2 := sc1.SubComponent2Pointer.Get(ctx2)
	s.Equal(newComponent, resolved2)
}

func (s *fieldSuite) TestUnresolvableDeferredPointerError() {
	workflowKey := definition.NewWorkflowKey(
		primitives.NewUUID().String(),
		primitives.NewUUID().String(),
		primitives.NewUUID().String(),
	)
	s.nodeBackend = &MockNodeBackend{
		HandleNextTransitionCount: func() int64 { return 1 },
		HandleGetCurrentVersion:   func() int64 { return 1 },
		HandleGetWorkflowKey:      func() definition.WorkflowKey { return workflowKey },
	}

	s.logger.(*testlogger.TestLogger).
		Expect(testlogger.Error, "failed to resolve deferred pointer during transaction close")

	orphanComponent := &TestSubComponent11{
		SubComponent11Data: &protoMessageType{
			CreateRequestId: "orphan-component",
		},
	}

	rootComponent := &TestComponent{
		ComponentData: &protoMessageType{
			CreateRequestId: "component-data",
		},
	}

	rootNode, ctx, err := s.setupComponentWithTree(rootComponent)
	s.NoError(err)

	// Get component from tree to mark node as needing sync.
	rootComponentInterface, err := rootNode.Component(ctx, ComponentRef{})
	s.NoError(err)
	rootComponent = rootComponentInterface.(*TestComponent)

	rootComponent.SubComponent11Pointer = ComponentPointerTo(ctx, orphanComponent)
	s.Equal(fieldTypeDeferredPointer, rootComponent.SubComponent11Pointer.Internal.fieldType())

	_, err = rootNode.CloseTransaction()
	s.Error(err)
	s.Contains(err.Error(), "failed to resolve deferred pointer during transaction close")
}