feat: add Mock method to ftl.Map (#1616)

Fixes: #1586

go-runtime/ftl/ftltest/fake.go

@@ -2,17 +2,28 @@ package ftltest
 
 import (
 	"context"
+	"fmt"
+	"reflect"
+	"strings"
 
 	"github.com/TBD54566975/ftl/go-runtime/internal"
 )
 
 type fakeFTL struct {
-	fsm *fakeFSMManager
+	fsm           *fakeFSMManager
+	mockMaps      map[uintptr]mapImpl
+	allowMapCalls bool
 }
 
+// mapImpl is a function that takes an object and returns an object of a potentially different
+// type but is not constrained by input/output type like ftl.Map.
+type mapImpl func(context.Context) (any, error)
+
 func newFakeFTL() *fakeFTL {
 	return &fakeFTL{
-		fsm: newFakeFSMManager(),
+		fsm:           newFakeFSMManager(),
+		mockMaps:      make(map[uintptr]mapImpl),
+		allowMapCalls: false,
 	}
 }
 
@@ -25,3 +36,48 @@ func (f *fakeFTL) FSMSend(ctx context.Context, fsm string, instance string, even
 func (f *fakeFTL) PublishEvent(ctx context.Context, topic string, event any) error {
 	panic("not implemented")
 }
+
+// addMapMock saves a new mock of ftl.Map to the internal map in fakeFTL.
+//
+// mockMap provides the whole mock implemention, so it gets called in place of both `fn`
+// and `getter` in ftl.Map.
+func (f *fakeFTL) addMapMock(mapper any, mockMap func(context.Context) (any, error)) {
+	key := makeMapKey(mapper)
+	f.mockMaps[key] = mockMap
+}
+
+func (f *fakeFTL) startAllowingMapCalls() {
+	f.allowMapCalls = true
+}
+
+func (f *fakeFTL) CallMap(ctx context.Context, mapper any, mapImpl func(context.Context) (any, error)) any {
+	key := makeMapKey(mapper)
+	mockMap, ok := f.mockMaps[key]
+	if ok {
+		return actuallyCallMap(ctx, mockMap)
+	}
+	if f.allowMapCalls {
+		return actuallyCallMap(ctx, mapImpl)
+	}
+	panic("map calls not allowed in tests by default. ftltest.Context should be instantiated with either ftltest.WithMapsAllowed() or a mock for the specific map being called using ftltest.WhenMap(...)")
+}
+
+func makeMapKey(mapper any) uintptr {
+	v := reflect.ValueOf(mapper)
+	if v.Kind() != reflect.Pointer {
+		panic("fakeFTL received object that was not a pointer, expected *MapHandle")
+	}
+	underlying := v.Elem().Type().Name()
+	if !strings.HasPrefix(underlying, "MapHandle[") {
+		panic(fmt.Sprintf("fakeFTL received *%s, expected *MapHandle", underlying))
+	}
+	return v.Pointer()
+}
+
+func actuallyCallMap(ctx context.Context, impl mapImpl) any {
+	out, err := impl(ctx)
+	if err != nil {
+		panic(err)
+	}
+	return out
+}

go-runtime/ftl/ftltest/ftltest.go

@@ -37,16 +37,17 @@ type Option func(context.Context, *OptionsState) error
 
 // Context suitable for use in testing FTL verbs with provided options
 func Context(options ...Option) context.Context {
-	ctx := log.ContextWithNewDefaultLogger(context.Background())
-	ctx = internal.WithContext(ctx, newFakeFTL())
-	name := reflection.Module()
-
 	state := &OptionsState{
 		configs:   make(map[string][]byte),
 		secrets:   make(map[string][]byte),
 		databases: make(map[string]modulecontext.Database),
 		mockVerbs: make(map[schema.RefKey]modulecontext.Verb),
 	}
+
+	ctx := log.ContextWithNewDefaultLogger(context.Background())
+	ctx = internal.WithContext(ctx, newFakeFTL())
+	name := reflection.Module()
+
 	for _, option := range options {
 		err := option(ctx, state)
 		if err != nil {
@@ -322,3 +323,41 @@ func WithCallsAllowedWithinModule() Option {
 		return nil
 	}
 }
+
+// WhenMap injects a fake implementation of a Mapping function
+//
+// To be used when setting up a context for a test:
+//
+//	ctx := ftltest.Context(
+//		ftltest.WhenMap(Example.MapHandle, func(ctx context.Context) (U, error) {
+//	    	// ...
+//		}),
+//		// ... other options
+//	)
+func WhenMap[T, U any](mapper *ftl.MapHandle[T, U], fake func(context.Context) (any, error)) Option {
+	return func(ctx context.Context, state *OptionsState) error {
+		someFTL := internal.FromContext(ctx)
+		fakeFTL, ok := someFTL.(*fakeFTL)
+		if !ok {
+			return fmt.Errorf("could not retrieve fakeFTL for saving a mock Map in test")
+		}
+		fakeFTL.addMapMock(mapper, fake)
+		return nil
+	}
+}
+
+// WithMapsAllowed allows all `ftl.Map` calls to pass through to their original
+// implementation.
+//
+// Any overrides provided by calling WhenMap(...) will take precedence.
+func WithMapsAllowed() Option {
+	return func(ctx context.Context, state *OptionsState) error {
+		someFTL := internal.FromContext(ctx)
+		fakeFTL, ok := someFTL.(*fakeFTL)
+		if !ok {
+			return fmt.Errorf("could not retrieve fakeFTL for saving a mock Map in test")
+		}
+		fakeFTL.startAllowingMapCalls()
+		return nil
+	}
+}

go-runtime/ftl/integration_test.go

@@ -0,0 +1,17 @@
+//go:build integration
+
+package ftl
+
+import (
+	"testing"
+
+	in "github.com/TBD54566975/ftl/integration"
+)
+
+func TestFTLMap(t *testing.T) {
+	in.Run(t, "",
+		in.CopyModule("mapper"),
+		in.Build("mapper"),
+		in.ExecModuleTest("mapper"),
+	)
+}

go-runtime/ftl/map.go

@@ -2,6 +2,9 @@ package ftl
 
 import (
 	"context"
+	"fmt"
+
+	"github.com/TBD54566975/ftl/go-runtime/internal"
 )
 
 type MapHandle[T, U any] struct {
@@ -10,14 +13,20 @@ type MapHandle[T, U any] struct {
 }
 
 func (mh *MapHandle[T, U]) Get(ctx context.Context) U {
-	t, err := mh.fn(ctx, mh.getter.Get(ctx))
-	if err != nil {
-		panic(err)
+	out := internal.FromContext(ctx).CallMap(ctx, mh, func(ctx context.Context) (any, error) {
+		return mh.fn(ctx, mh.getter.Get(ctx))
+	})
+	u, ok := out.(U)
+	if !ok {
+		panic(fmt.Sprintf("output object %v is not compatible with expected type %T", out, *new(U)))
 	}
-	return t
+	return u
 }
 
 // Map an FTL resource type to a new type.
 func Map[T, U any](getter Handle[T], fn func(context.Context, T) (U, error)) MapHandle[T, U] {
-	return MapHandle[T, U]{fn: fn, getter: getter}
+	return MapHandle[T, U]{
+		fn:     fn,
+		getter: getter,
+	}
 }

go-runtime/ftl/map_test.go

@@ -3,8 +3,10 @@ package ftl
 import (
 	"context"
 	"fmt"
+	"strconv"
 	"testing"
 
+	"github.com/TBD54566975/ftl/go-runtime/internal"
 	"github.com/alecthomas/assert/v2"
 )
 
@@ -13,7 +15,7 @@ type intHandle int
 func (s intHandle) Get(ctx context.Context) int { return int(s) }
 
 func TestMapPanic(t *testing.T) {
-	ctx := context.Background()
+	ctx := internal.WithContext(context.Background(), internal.New())
 	n := intHandle(1)
 	once := Map(n, func(ctx context.Context, n int) (string, error) {
 		return "", fmt.Errorf("test error %d", n)
@@ -22,3 +24,12 @@ func TestMapPanic(t *testing.T) {
 		once.Get(ctx)
 	})
 }
+
+func TestMapGet(t *testing.T) {
+	ctx := internal.WithContext(context.Background(), internal.New())
+	n := intHandle(1)
+	once := Map(n, func(ctx context.Context, n int) (string, error) {
+		return strconv.Itoa(n), nil
+	})
+	assert.Equal(t, once.Get(ctx), "1")
+}

go-runtime/ftl/testdata/go/mapper/ftl.toml

@@ -0,0 +1,2 @@
+module = "mapper"
+language = "go"

go-runtime/ftl/testdata/go/mapper/go.mod

@@ -0,0 +1,45 @@
+module ftl/mapper
+
+go 1.22.2
+
+toolchain go1.22.3
+
+require github.com/TBD54566975/ftl v1.1.0
+
+require (
+	connectrpc.com/connect v1.16.1 // indirect
+	connectrpc.com/grpcreflect v1.2.0 // indirect
+	connectrpc.com/otelconnect v0.7.0 // indirect
+	github.com/alecthomas/concurrency v0.0.2 // indirect
+	github.com/alecthomas/participle/v2 v2.1.1 // indirect
+	github.com/alecthomas/types v0.16.0 // indirect
+	github.com/alessio/shellescape v1.4.2 // indirect
+	github.com/danieljoos/wincred v1.2.0 // indirect
+	github.com/go-logr/logr v1.4.2 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
+	github.com/godbus/dbus/v5 v5.1.0 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+	github.com/jackc/pgx/v5 v5.6.0 // indirect
+	github.com/jackc/puddle/v2 v2.2.1 // indirect
+	github.com/jpillora/backoff v1.0.0 // indirect
+	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/multiformats/go-base36 v0.2.0 // indirect
+	github.com/swaggest/jsonschema-go v0.3.70 // indirect
+	github.com/swaggest/refl v1.3.0 // indirect
+	github.com/zalando/go-keyring v0.2.4 // indirect
+	go.opentelemetry.io/otel v1.27.0 // indirect
+	go.opentelemetry.io/otel/metric v1.27.0 // indirect
+	go.opentelemetry.io/otel/trace v1.27.0 // indirect
+	golang.org/x/crypto v0.23.0 // indirect
+	golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 // indirect
+	golang.org/x/mod v0.17.0 // indirect
+	golang.org/x/net v0.25.0 // indirect
+	golang.org/x/sync v0.7.0 // indirect
+	golang.org/x/sys v0.20.0 // indirect
+	golang.org/x/text v0.15.0 // indirect
+	google.golang.org/protobuf v1.34.1 // indirect
+)
+
+replace github.com/TBD54566975/ftl => ../../../../..