Testing and Mocking Plugins

Testing custom plugin logic within the Flyte execution framework requires isolating your code from external services like the Flyte Admin events service or the Catalog cache. By using the provided test doubles in the executor/pkg/controller package, you can verify that your plugins emit the correct events, handle abort signals properly, and interact correctly with the cache.

In this tutorial, you will build a test suite that validates a plugin's lifecycle using Flyte's internal mocking utilities.

Prerequisites

To follow this tutorial, you need the Flyte codebase and a Go environment configured. Your tests will primarily interact with classes found in:

• executor/pkg/controller
• github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core

Step 1: Verify Event Emission with recordingEventsClient

When a TaskAction transitions through different phases, the Flyte controller emits events to the Flyte Admin service. To verify these events without a running Admin instance, use the recordingEventsClient. This client captures all ActionEvent objects in a thread-safe slice.

import (
    "context"
    "github.com/flyteorg/flyte/v2/executor/pkg/controller"
    "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/common"
    "github.com/stretchr/testify/assert"
)

func TestPluginEventEmission(t *testing.T) {
    ctx := context.Background()
    
    // Initialize the recording client
    recorder := &recordingEventsClient{}
    
    // Inject it into the TaskActionReconciler
    reconciler := &TaskActionReconciler{
        // ... other fields like Client, DataStore, PluginRegistry ...
        eventsClient: recorder,
    }

    // Trigger a reconciliation that should result in an abort
    _, err := reconciler.Reconcile(ctx, reconcile.Request{NamespacedName: taskName})
    assert.NoError(t, err)

    // Retrieve and assert on recorded events
    recorded := recorder.RecordedEvents()
    assert.NotEmpty(t, recorded)
    
    foundAbort := false
    for _, event := range recorded {
        if event.GetPhase() == common.ActionPhase_ACTION_PHASE_ABORTED {
            foundAbort = true
            break
        }
    }
    assert.True(t, foundAbort, "Expected an ABORTED event to be recorded")
}

The recordingEventsClient uses a sync.Mutex internally, making it safe for use in concurrent reconciliation tests where multiple events might be recorded simultaneously.

Step 2: Mock Plugin Behavior with fakePlugin

To test how the controller interacts with your plugin's lifecycle methods (like Abort or Finalize), use the fakePlugin. This implementation allows you to track the number of calls made to specific methods.

func TestPluginAbortLogic(t *testing.T) {
    // Create a fake plugin with a specific ID
    myPlugin := &fakePlugin{
        id: "my-custom-plugin",
    }

    // Simulate an abort call from the controller
    err := myPlugin.Abort(ctx, taskCtx)
    assert.NoError(t, err)
    
    // Verify the plugin tracked the call
    assert.Equal(t, 1, myPlugin.abortCalls)
}

Use fakePlugin when you need to verify behavioral side effects. If you only need a structural implementation of the pluginsCore.Plugin interface for validation tests where the methods aren't actually called, use mockPlugin instead.

Step 3: Simulate Cache Hits and Misses with stubCatalogClient

The stubCatalogClient is a flexible mock for the catalog.Client interface. It uses function fields to allow you to define custom behavior for specific catalog operations like Get, Put, and reservation management.

import (
    "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/catalog"
    "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/io"
    corepb "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
)

func TestCacheHitScenario(t *testing.T) {
    // Define a stub that simulates a cache hit
    catalogStub := &stubCatalogClient{
        getFunc: func(ctx context.Context, key catalog.Key) (catalog.Entry, error) {
            return catalog.NewCatalogEntry(
                ioutils.NewInMemoryOutputReader(expectedOutputs, nil, nil),
                catalog.NewStatus(corepb.CatalogCacheStatus_CACHE_HIT, nil),
            ), nil
        },
    }

    reconciler := &TaskActionReconciler{
        Catalog: catalogStub,
        // ... other fields ...
    }

    // Evaluate cache before execution
    transition, handled, err := reconciler.evaluateCacheBeforeExecution(ctx, taskAction, tCtx)
    assert.NoError(t, err)
    assert.True(t, handled)
    assert.Equal(t, pluginsCore.PhaseSuccess, transition.Info().Phase())
}

Warning: You must initialize the specific function fields (like getFunc or putFunc) that your code path will invoke. If a method is called on stubCatalogClient and its corresponding function field is nil, the test will panic.

Step 4: Test Resolution Failures with mockPluginResolver

When testing validation logic, you often need to simulate scenarios where a plugin cannot be found for a specific task type. The mockPluginResolver allows you to return a pre-configured error.

func TestValidateTaskAction_PluginNotFound(t *testing.T) {
    // Configure the resolver to return an error
    resolver := &mockPluginResolver{
        plugin: nil,
        err:    fmt.Errorf("no plugin registered for task type %q", "container"),
    }

    // Validate a TaskAction spec against this resolver
    _, reason, err := validateTaskAction(validTaskAction(), resolver)
    
    assert.Error(t, err)
    assert.Equal(t, flyteorgv1.ConditionReasonPluginNotFound, reason)
}

Summary

By combining these mocks, you can create comprehensive unit and integration tests for Flyte plugins:

• Use recordingEventsClient to verify the sequence of ActionEvent emissions.
• Use fakePlugin to track calls to Abort and Finalize.
• Use stubCatalogClient to control cache behavior (hits, misses, and errors).
• Use mockPluginResolver to test edge cases in plugin discovery and validation.

For more complex scenarios, you can combine these with fakeEventsClient, which provides a no-op implementation of the events service when event recording is not the focus of your test.