Testing and Mocking Services

Flyte uses interfaces and generated mocks to isolate service logic from infrastructure dependencies like Kubernetes, Cloud Storage, and external APIs. This allows you to write fast, deterministic unit tests for core logic without requiring a live cluster.

Mocking High-Level Services

For high-level services like the Actions Service, Flyte uses mockery to generate type-safe mocks. These mocks are typically found in a mocks sub-package relative to the interface definition.

When testing a service, you instantiate the mock using the New[Interface](t) constructor. This pattern ensures that the mock is automatically cleaned up and its expectations are verified when the test completes.

// actions/service/actions_service_test.go

func TestEnqueue(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		// Instantiate the mock with the testing object for auto-cleanup
		m := mocks.NewActionsClientInterface(t)
		svc := NewActionsService(m)

		// Set expectations using the type-safe EXPECT() method
		m.EXPECT().Enqueue(mock.Anything, testAction, (*task.RunSpec)(nil)).Return(nil)

		resp, err := svc.Enqueue(context.Background(), connect.NewRequest(&actions.EnqueueRequest{
			Action: testAction,
		}))

		assert.NoError(t, err)
		assert.NotNil(t, resp)
	})
}

Testing Task Plugins

Testing Flyte task plugins requires mocking complex execution contexts like TaskExecutionContext and SetupContext. These interfaces provide access to input readers, task templates, and metadata.

Mocking TaskExecutionContext

The TaskExecutionContext is a composite interface. To mock it effectively, you often need to mock its constituent parts such as TaskReader and InputReader.

// flyteplugins/go/tasks/plugins/core/sleep/plugin_test.go

func newTaskExecutionContext(sleepDuration time.Duration, generatedName string) *coreMocks.TaskExecutionContext {
	// Mock the TaskReader to return a specific template
	taskReader := &coreMocks.TaskReader{}
	taskReader.EXPECT().Read(mock.Anything).Return(taskTemplate, nil)

	// Mock the InputReader to return specific literals
	inputReader := &ioMocks.InputReader{}
	inputReader.EXPECT().Get(mock.Anything).Return(inputs, nil)

	// Mock metadata and IDs
	taskExecutionID := &coreMocks.TaskExecutionID{}
	taskExecutionID.EXPECT().GetGeneratedName().Return(generatedName)

	metadata := &coreMocks.TaskExecutionMetadata{}
	metadata.EXPECT().GetTaskExecutionID().Return(taskExecutionID)

	// Compose the main context mock
	tCtx := &coreMocks.TaskExecutionContext{}
	tCtx.EXPECT().TaskReader().Return(taskReader)
	tCtx.EXPECT().InputReader().Return(inputReader)
	tCtx.EXPECT().TaskExecutionMetadata().Return(metadata)

	return tCtx
}

Mocking SetupContext

When testing plugin loading or initialization, use the SetupContext mock.

// flyteplugins/go/tasks/pluginmachinery/core/plugin_test.go

func TestLoadPlugin(t *testing.T) {
	t.Run("valid", func(t *testing.T) {
		corePlugin := &mocks.Plugin{}
		corePlugin.EXPECT().GetID().Return("core")
		corePlugin.EXPECT().GetProperties().Return(core.PluginProperties{})

		setupCtx := mocks.SetupContext{}
		p, err := core.LoadPlugin(context.TODO(), &setupCtx, corePluginEntry)
		
		assert.NilError(t, err)
		assert.Equal(t, "core", p.GetID())
	})
}

Mocking Storage Interactions

Flyte provides a RawStore interface in flytestdlib for interacting with storage backends (S3, GCS, Local). You can use the generated RawStore mock to verify data interactions.

// flytestdlib/storage/mocks/mocks.go

func TestStorageUsage(t *testing.T) {
    m := &mocks.RawStore{}
    m.EXPECT().CopyRaw(mock.Anything, source, dest, opts).Return(nil)
    
    err := myService.MoveData(ctx, source, dest)
    assert.NoError(t, err)
}

For lower-level testing of the storage implementation itself, Flyte uses hand-written mocks for the stow library to simulate different cloud provider behaviors.

// flytestdlib/storage/stow_store_test.go

type mockStowLoc struct {
	stow.Location
	ContainerCb       func(id string) (stow.Container, error)
	CreateContainerCb func(name string) (stow.Container, error)
}

func (m mockStowLoc) Container(id string) (stow.Container, error) {
	return m.ContainerCb(id)
}

Best Practices

Automatic Expectation Verification

Always prefer the New[Interface](t) constructor when available. It registers a cleanup function on *testing.T that automatically calls AssertExpectations(t) at the end of the test.

Type-Safe Expectations

Use the .EXPECT() method provided by newer mockery templates. It provides a fluent, type-safe API for setting up mocks, which helps catch signature changes at compile time rather than runtime.

Handling Channels

When mocking methods that return or interact with channels (such as subscription services), ensure you manage the channel lifecycle within the test. Failing to close channels or provide a consumer can lead to goroutine leaks or deadlocks in the test suite.

Mocking External Errors

To test error handling, use the .Return() method to return specific error types, such as connect.CodeInternal for RPC services or stow.ErrNotFound for storage.

m.EXPECT().Enqueue(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("k8s error"))