From 8245b90294e64edfb79728647ffe3d575f3d5943 Mon Sep 17 00:00:00 2001
From: Matt Toohey <mtoohey@block.xyz>
Date: Fri, 14 Jun 2024 12:27:09 +1000
Subject: [PATCH] fix: WhenMap() now typesafe, ftl.Map returns ptr (#1777)

- `ftltest.WhenMap(...)` included a func that returned `any` when it
should be strongly types with the map's return type
- `ftl.Map` was returning `MapHandle` where as `handle.Get()` was
defined on `*MapHandle` and `ftltest.WhenMap(...)` expected a pointer as
well. I've made it all consistently use `*MapHandle` so you don't need
to do `&handle` anywhere
---
 go-runtime/ftl/ftltest/fake.go                   | 7 +++++--
 go-runtime/ftl/ftltest/ftltest.go                | 4 ++--
 go-runtime/ftl/map.go                            | 4 ++--
 go-runtime/ftl/testdata/go/mapper/mapper_test.go | 2 +-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/go-runtime/ftl/ftltest/fake.go b/go-runtime/ftl/ftltest/fake.go
index e3ebe1d5d0..efd35b323f 100644
--- a/go-runtime/ftl/ftltest/fake.go
+++ b/go-runtime/ftl/ftltest/fake.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/TBD54566975/ftl/backend/schema"
 	"github.com/TBD54566975/ftl/common/configuration"
+	"github.com/TBD54566975/ftl/go-runtime/ftl"
 	"github.com/TBD54566975/ftl/go-runtime/internal"
 	"github.com/alecthomas/types/optional"
 )
@@ -119,9 +120,11 @@ func (f *fakeFTL) FSMSend(ctx context.Context, fsm string, instance string, even
 //
 // 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)) {
+func addMapMock[T, U any](f *fakeFTL, mapper *ftl.MapHandle[T, U], mockMap func(context.Context) (U, error)) {
 	key := makeMapKey(mapper)
-	f.mockMaps[key] = mockMap
+	f.mockMaps[key] = func(ctx context.Context) (any, error) {
+		return mockMap(ctx)
+	}
 }
 
 func (f *fakeFTL) startAllowingMapCalls() {
diff --git a/go-runtime/ftl/ftltest/ftltest.go b/go-runtime/ftl/ftltest/ftltest.go
index a5978b4ae6..b6e75a7bc7 100644
--- a/go-runtime/ftl/ftltest/ftltest.go
+++ b/go-runtime/ftl/ftltest/ftltest.go
@@ -332,10 +332,10 @@ func WithCallsAllowedWithinModule() Option {
 //		}),
 //		// ... other options
 //	)
-func WhenMap[T, U any](mapper *ftl.MapHandle[T, U], fake func(context.Context) (any, error)) Option {
+func WhenMap[T, U any](mapper *ftl.MapHandle[T, U], fake func(context.Context) (U, error)) Option {
 	return func(ctx context.Context, state *OptionsState) error {
 		fftl := internal.FromContext(ctx).(*fakeFTL) //nolint:forcetypeassert
-		fftl.addMapMock(mapper, fake)
+		addMapMock(fftl, mapper, fake)
 		return nil
 	}
 }
diff --git a/go-runtime/ftl/map.go b/go-runtime/ftl/map.go
index fe1de88a78..9b1604b9a5 100644
--- a/go-runtime/ftl/map.go
+++ b/go-runtime/ftl/map.go
@@ -26,8 +26,8 @@ func (mh *MapHandle[T, U]) Get(ctx context.Context) 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]{
+func Map[T, U any](getter Handle[T], fn func(context.Context, T) (U, error)) *MapHandle[T, U] {
+	return &MapHandle[T, U]{
 		fn:     fn,
 		handle: getter,
 	}
diff --git a/go-runtime/ftl/testdata/go/mapper/mapper_test.go b/go-runtime/ftl/testdata/go/mapper/mapper_test.go
index 09d4c04409..8136b62db4 100644
--- a/go-runtime/ftl/testdata/go/mapper/mapper_test.go
+++ b/go-runtime/ftl/testdata/go/mapper/mapper_test.go
@@ -23,7 +23,7 @@ func TestPanicsWithoutExplicitlyAllowingMaps(t *testing.T) {
 
 func TestMockGet(t *testing.T) {
 	mockOut := 123
-	ctx := ftltest.Context(ftltest.WhenMap(&m, func(ctx context.Context) (any, error) {
+	ctx := ftltest.Context(ftltest.WhenMap(m, func(ctx context.Context) (int, error) {
 		return mockOut, nil
 	}))
 	got := m.Get(ctx)