diff --git a/go-runtime/ftl/ftltest/fake.go b/go-runtime/ftl/ftltest/fake.go
index fc2b5c8301..c94562c9d7 100644
--- a/go-runtime/ftl/ftltest/fake.go
+++ b/go-runtime/ftl/ftltest/fake.go
@@ -64,17 +64,17 @@ type fakeFTL struct {
 // type but is not constrained by input/output type like ftl.Map.
 type mapImpl func(context.Context) (any, error)
 
-func newFakeFTL(ctx context.Context) *fakeFTL {
+func contextWithFakeFTL(ctx context.Context) context.Context {
 	fake := &fakeFTL{
 		fsm:           newFakeFSMManager(),
 		mockMaps:      map[uintptr]mapImpl{},
 		allowMapCalls: false,
 		configValues:  map[string][]byte{},
 		secretValues:  map[string][]byte{},
-		pubSub:        newFakePubSub(ctx),
 	}
-
-	return fake
+	ctx = internal.WithContext(ctx, fake)
+	fake.pubSub = newFakePubSub(ctx)
+	return ctx
 }
 
 var _ internal.FTL = &fakeFTL{}
diff --git a/go-runtime/ftl/ftltest/ftltest.go b/go-runtime/ftl/ftltest/ftltest.go
index f61cd09f37..a39487a5df 100644
--- a/go-runtime/ftl/ftltest/ftltest.go
+++ b/go-runtime/ftl/ftltest/ftltest.go
@@ -51,7 +51,7 @@ func Context(options ...Option) context.Context {
 	}
 
 	ctx := log.ContextWithNewDefaultLogger(context.Background())
-	ctx = internal.WithContext(ctx, newFakeFTL(ctx))
+	ctx = contextWithFakeFTL(ctx)
 	name := reflection.Module()
 
 	sort.Slice(options, func(i, j int) bool {
diff --git a/go-runtime/ftl/ftltest/ftltest_test.go b/go-runtime/ftl/ftltest/ftltest_test.go
index 41d192a596..cc9a5420a3 100644
--- a/go-runtime/ftl/ftltest/ftltest_test.go
+++ b/go-runtime/ftl/ftltest/ftltest_test.go
@@ -6,7 +6,6 @@ import (
 	"testing"
 
 	"github.com/TBD54566975/ftl/go-runtime/ftl"
-	"github.com/TBD54566975/ftl/go-runtime/internal"
 	"github.com/TBD54566975/ftl/internal/log"
 	"github.com/alecthomas/assert/v2"
 )
@@ -27,7 +26,7 @@ func PanicsWithErr(t testing.TB, substr string, fn func()) {
 
 func TestFtlTestProjectNotLoadedInContext(t *testing.T) {
 	ctx := log.ContextWithNewDefaultLogger(context.Background())
-	ctx = internal.WithContext(ctx, newFakeFTL(ctx))
+	ctx = contextWithFakeFTL(ctx)
 
 	// This should panic suggesting to use ftltest.WithDefaultProjectFile()
 	PanicsWithErr(t, "ftltest.WithDefaultProjectFile()", func() {