Automatically stop server after test completion via T.Cleanup

temporaltest/options.go

@@ -11,6 +11,9 @@ type TestServerOption interface {
 }
 
 // WithT directs all worker and client logs to the test logger.
+//
+// If this option is specified, then server will automatically be stopped when the
+// test completes.
 func WithT(t *testing.T) TestServerOption {
 	return newApplyFuncContainer(func(server *TestServer) {
 		server.t = t

temporaltest/server.go

@@ -96,8 +96,10 @@ func (ts *TestServer) Stop() {
 	ts.server.Stop()
 }
 
-// NewServer starts and returns a new TestServer. The caller should call Stop
-// when finished, to shut it down.
+// NewServer starts and returns a new TestServer.
+//
+// If not specifying the WithT option, the caller should execute Stop when finished to close
+// the server and release resources.
 func NewServer(opts ...TestServerOption) *TestServer {
 	rand.Seed(time.Now().UnixNano())
 	testNamespace := fmt.Sprintf("temporaltest-%d", rand.Intn(999999))
@@ -111,6 +113,12 @@ func NewServer(opts ...TestServerOption) *TestServer {
 		opt.apply(&ts)
 	}
 
+	if ts.t != nil {
+		ts.t.Cleanup(func() {
+			ts.Stop()
+		})
+	}
+
 	s, err := temporalite.NewServer(
 		temporalite.WithNamespaces(ts.defaultTestNamespace),
 		temporalite.WithPersistenceDisabled(),

temporaltest/server_test.go

@@ -10,10 +10,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/DataDog/temporalite/internal/examples/helloworld"
-	"github.com/DataDog/temporalite/temporaltest"
 	"go.temporal.io/sdk/client"
 	"go.temporal.io/sdk/worker"
+
+	"github.com/DataDog/temporalite/internal/examples/helloworld"
+	"github.com/DataDog/temporalite/temporaltest"
 )
 
 // to be used in example code
@@ -22,17 +23,13 @@ var t *testing.T
 func ExampleNewServer_testWorker() {
 	// Create test Temporal server and client
 	ts := temporaltest.NewServer(temporaltest.WithT(t))
-	// Stop server and close clients when tests complete
-	defer ts.Stop()
+	c := ts.Client()
 
 	// Register a new worker on the `hello_world` task queue
 	ts.Worker("hello_world", func(registry worker.Registry) {
 		helloworld.RegisterWorkflowsAndActivities(registry)
 	})
 
-	// Create a test client
-	c := ts.Client()
-
 	// Start test workflow
 	wfr, err := c.ExecuteWorkflow(
 		context.Background(),
@@ -57,7 +54,6 @@ func ExampleNewServer_testWorker() {
 
 func TestNewServer(t *testing.T) {
 	ts := temporaltest.NewServer(temporaltest.WithT(t))
-	defer ts.Stop()
 
 	ts.Worker("hello_world", func(registry worker.Registry) {
 		helloworld.RegisterWorkflowsAndActivities(registry)