rip out last bit of modulecontext changes

go-runtime/ftl/ftltest/fake.go

@@ -7,12 +7,12 @@ import (
 	"strings"
 
 	"github.com/TBD54566975/ftl/go-runtime/internal"
-	"github.com/TBD54566975/ftl/internal/modulecontext"
 )
 
 type fakeFTL struct {
-	fsm      *fakeFSMManager
-	mockMaps map[uintptr]mapImpl
+	fsm           *fakeFSMManager
+	mockMaps      map[uintptr]mapImpl
+	allowMapCalls bool
 }
 
 // mapImpl is a function that takes an object and returns an object of a potentially different
@@ -21,8 +21,9 @@ type mapImpl func(context.Context) (any, error)
 
 func newFakeFTL() *fakeFTL {
 	return &fakeFTL{
-		fsm:      newFakeFSMManager(),
-		mockMaps: make(map[uintptr]mapImpl),
+		fsm:           newFakeFSMManager(),
+		mockMaps:      make(map[uintptr]mapImpl),
+		allowMapCalls: false,
 	}
 }
 
@@ -41,13 +42,17 @@ func (f *fakeFTL) addMapMock(ctx context.Context, mapper any, mockMap func(conte
 	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 modulecontext.FromContext(ctx).AllowDirectMapBehavior() {
+	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(...)")

go-runtime/ftl/ftltest/ftltest.go

@@ -31,7 +31,6 @@ type OptionsState struct {
 	databases               map[string]modulecontext.Database
 	mockVerbs               map[schema.RefKey]modulecontext.Verb
 	allowDirectVerbBehavior bool
-	allowDirectMapBehavior  bool
 }
 
 type Option func(context.Context, *OptionsState) error
@@ -57,7 +56,7 @@ func Context(options ...Option) context.Context {
 	}
 
 	builder := modulecontext.NewBuilder(name).AddConfigs(state.configs).AddSecrets(state.secrets).AddDatabases(state.databases)
-	builder = builder.UpdateForTesting(state.mockVerbs, state.allowDirectVerbBehavior, state.allowDirectMapBehavior, newFakeLeaseClient())
+	builder = builder.UpdateForTesting(state.mockVerbs, state.allowDirectVerbBehavior, newFakeLeaseClient())
 	return builder.Build().ApplyToContext(ctx)
 }
 
@@ -353,7 +352,12 @@ func WhenMap[T, U any](mapper *ftl.MapHandle[T, U], fake func(context.Context) (
 // Any overrides provided by calling WhenMap(...) will take precedence.
 func WithMapsAllowed() Option {
 	return func(ctx context.Context, state *OptionsState) error {
-		state.allowDirectMapBehavior = true
+		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
 	}
 }

internal/modulecontext/module_context.go

@@ -32,7 +32,6 @@ type ModuleContext struct {
 	isTesting               bool
 	mockVerbs               map[schema.RefKey]Verb
 	allowDirectVerbBehavior bool
-	allowDirectMapBehavior  bool
 	leaseClient             optional.Option[LeaseClient]
 }
 
@@ -77,21 +76,12 @@ func (b *Builder) AddDatabases(databases map[string]Database) *Builder {
 }
 
 // UpdateForTesting marks the builder as part of a test environment and adds mock verbs and flags for other test features.
-func (b *Builder) UpdateForTesting(
-	mockVerbs map[schema.RefKey]Verb,
-	allowDirectVerbBehavior bool,
-	allowDirectMapBehavior bool,
-	leaseClient LeaseClient,
-) *Builder {
+func (b *Builder) UpdateForTesting(mockVerbs map[schema.RefKey]Verb, allowDirectVerbBehavior bool, leaseClient LeaseClient) *Builder {
 	b.isTesting = true
-
 	for name, verb := range mockVerbs {
 		b.mockVerbs[name] = verb
 	}
 	b.allowDirectVerbBehavior = allowDirectVerbBehavior
-
-	b.allowDirectMapBehavior = allowDirectMapBehavior
-
 	b.leaseClient = optional.Some[LeaseClient](leaseClient)
 	return b
 }
@@ -167,11 +157,6 @@ func (m ModuleContext) MockLeaseClient() optional.Option[LeaseClient] {
 	return m.leaseClient
 }
 
-// AllowDirectMapBehavior returns whether ftl.Map calls should be allowed in test
-func (m ModuleContext) AllowDirectMapBehavior() bool {
-	return m.allowDirectMapBehavior
-}
-
 // BehaviorForVerb returns what to do to execute a verb
 //
 // This allows module context to dictate behavior based on testing options