diff --git a/backend/runner/pubsub/testdata/go/publisher/types.ftl.go b/backend/runner/pubsub/testdata/go/publisher/types.ftl.go
index 0bd321bb3..be317cc55 100644
--- a/backend/runner/pubsub/testdata/go/publisher/types.ftl.go
+++ b/backend/runner/pubsub/testdata/go/publisher/types.ftl.go
@@ -3,8 +3,8 @@ package publisher
 
 import (
 	"context"
-	"github.com/block/ftl/go-runtime/ftl"
 	"github.com/block/ftl/common/reflection"
+	"github.com/block/ftl/go-runtime/ftl"
 	"github.com/block/ftl/go-runtime/server"
 )
 
@@ -18,7 +18,7 @@ func init() {
 	reflection.Register(
 		reflection.ProvideResourcesForVerb(
 			PublishOne,
-			server.TopicHandle[PubSubEvent, ftl.SinglePartitionMap[PubSubEvent]]("publisher", "testTopic"),
+			server.TopicHandle[PubSubEvent, PartitionMapper]("publisher", "testTopic"),
 		),
 		reflection.ProvideResourcesForVerb(
 			PublishOneToTopic2,
@@ -26,7 +26,7 @@ func init() {
 		),
 		reflection.ProvideResourcesForVerb(
 			PublishTen,
-			server.TopicHandle[PubSubEvent, ftl.SinglePartitionMap[PubSubEvent]]("publisher", "testTopic"),
+			server.TopicHandle[PubSubEvent, PartitionMapper]("publisher", "testTopic"),
 		),
 	)
 }
diff --git a/backend/runner/pubsub/testdata/java/subscriber/src/main/java/xyz/block/ftl/java/test/subscriber/Subscriber.java b/backend/runner/pubsub/testdata/java/subscriber/src/main/java/xyz/block/ftl/java/test/subscriber/Subscriber.java
index 2148cb360..ca74499b6 100644
--- a/backend/runner/pubsub/testdata/java/subscriber/src/main/java/xyz/block/ftl/java/test/subscriber/Subscriber.java
+++ b/backend/runner/pubsub/testdata/java/subscriber/src/main/java/xyz/block/ftl/java/test/subscriber/Subscriber.java
@@ -1,7 +1,5 @@
 package xyz.block.ftl.java.test.subscriber;
 
-import java.util.concurrent.atomic.AtomicInteger;
-
 import ftl.publisher.PubSubEvent;
 import ftl.publisher.TestTopicTopic;
 import ftl.publisher.Topic2Topic;