pubsub-go/jetstream/test/publish_test.go

// Copyright 2022-2025 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package test

import (
	"context"
	"errors"
	"fmt"
	"os"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/hanzoai/pubsub-go"
	"github.com/hanzoai/pubsub-go/jetstream"
	"github.com/nats-io/nats-server/v2/server"
)

func TestPublishMsg(t *testing.T) {
	type publishConfig struct {
		msg             *nats.Msg
		opts            []jetstream.PublishOpt
		expectedHeaders nats.Header
		expectedAck     jetstream.PubAck
		withError       func(*testing.T, error)
	}
	tests := []struct {
		name      string
		srvConfig []byte
		timeout   time.Duration
		msgs      []publishConfig
	}{
		{
			name: "publish 3 simple messages, no opts",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
						Domain:   "",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
						Domain:   "",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
						Domain:   "",
					},
				},
			},
		},
		{
			name: "publish 3 messages with message ID, with duplicate",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:    "foo",
						Sequence:  1,
						Duplicate: true,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"2"},
					},
				},
			},
		},
		{
			name: "expect last msg ID",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"2"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastMsgID("2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Msg-Id": []string{"2"},
					},
				},
			},
		},
		{
			name: "invalid last msg id",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastMsgID("abc")},
					withError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10070 {
							t.Fatalf("Expected error code: 10070; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect last sequence",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequence(2)},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Sequence": []string{"2"},
					},
				},
			},
		},
		{
			name: "invalid last seq",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequence(123)},
					withError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10071 {
							t.Fatalf("Expected error code: 10071; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect last sequence per subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequencePerSubject(1)},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Subject-Sequence": []string{"1"},
					},
				},
			},
		},
		{
			name: "invalid last sequence per subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequencePerSubject(123)},
					withError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10071 {
							t.Fatalf("Expected error code: 10071; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect last sequence for subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequenceForSubject(2, "FOO.2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Subject-Sequence":         []string{"2"},
						"Nats-Expected-Last-Subject-Sequence-Subject": []string{"FOO.2"},
					},
				},
			},
		},
		{
			name: "invalid last sequence for subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequenceForSubject(123, "FOO.2")},
					withError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10071 {
							t.Fatalf("Expected error code: 10071; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect stream header",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectStream("foo")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Stream": []string{"foo"},
					},
				},
			},
		},
		{
			name: "invalid expected stream header",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectStream("abc")},
					withError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10060 {
							t.Fatalf("Expected error code: 10060; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "publish 3 simple messages with domain set",
			srvConfig: []byte(
				`
					listen: 127.0.0.1:-1
					jetstream: {domain: "test-domain"}
				`),
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
						Domain:   "test-domain",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
						Domain:   "test-domain",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
						Domain:   "test-domain",
					},
				},
			},
		},
		{
			name: "publish timeout",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Stream": []string{"foo"},
					},
					withError: func(t *testing.T, err error) {
						if !errors.Is(err, context.DeadlineExceeded) {
							t.Fatalf("Expected deadline exceeded error; got: %v", err)
						}
					},
				},
			},
			timeout: 1 * time.Nanosecond,
		},
		{
			name: "invalid option set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithStallWait(1 * time.Second)},
					withError: func(t *testing.T, err error) {
						if !errors.Is(err, jetstream.ErrInvalidOption) {
							t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
						}
					},
				},
			},
		},
		{
			name: "no subject set on message",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data: []byte("msg 1"),
					},
					opts: []jetstream.PublishOpt{},
					withError: func(t *testing.T, err error) {
						if !errors.Is(err, nats.ErrBadSubject) {
							t.Fatalf("Expected error: %v; got: %v", nats.ErrBadSubject, err)
						}
					},
				},
			},
		},
		{
			name: "invalid subject set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "ABC",
					},
					opts: []jetstream.PublishOpt{},
					withError: func(t *testing.T, err error) {
						if !errors.Is(err, jetstream.ErrNoStreamResponse) {
							t.Fatalf("Expected error: %v; got: %v", jetstream.ErrNoStreamResponse, err)
						}
					},
				},
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var srv *server.Server
			if test.srvConfig != nil {
				conf := createConfFile(t, test.srvConfig)
				defer os.Remove(conf)
				srv, _ = RunServerWithConfig(conf)
			} else {
				srv = RunBasicJetStreamServer()
			}
			defer shutdownJSServerAndRemoveStorage(t, srv)
			nc, err := nats.Connect(srv.ClientURL())
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}

			js, err := jetstream.New(nc)
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			defer nc.Close()

			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}, MaxMsgSize: 128})
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}

			for _, pub := range test.msgs {
				var pubCtx context.Context
				var pubCancel context.CancelFunc
				if test.timeout != 0 {
					pubCtx, pubCancel = context.WithTimeout(ctx, test.timeout)
				} else {
					pubCtx, pubCancel = context.WithTimeout(ctx, 1*time.Minute)
				}
				ack, err := js.PublishMsg(pubCtx, pub.msg, pub.opts...)
				pubCancel()
				if pub.withError != nil {
					pub.withError(t, err)
					continue
				}
				if err != nil {
					t.Fatalf("Unexpected error: %v", err)
				}
				if !reflect.DeepEqual(pub.expectedHeaders, pub.msg.Header) {
					t.Fatalf("Invalid headers on message; want: %v; got: %v", pub.expectedHeaders, pub.msg.Header)
				}
				if *ack != pub.expectedAck {
					t.Fatalf("Invalid ack received; want: %v; got: %v", pub.expectedAck, ack)
				}
			}
		})
	}
}

func TestPublishWithTTL(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name: "foo", Subjects: []string{"FOO.*"}, MaxMsgSize: 64, AllowMsgTTL: true})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ack, err := js.Publish(ctx, "FOO.1", []byte("msg"), jetstream.WithMsgTTL(1*time.Second))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if ttl := gotMsg.Header.Get("Nats-TTL"); ttl != "1s" {
		t.Fatalf("Expected message to have TTL header set to 1s; got: %s", ttl)
	}
	time.Sleep(1500 * time.Millisecond)
	_, err = stream.GetMsg(ctx, ack.Sequence)
	if !errors.Is(err, jetstream.ErrMsgNotFound) {
		t.Fatalf("Expected not found error; got: %v", err)
	}
}

func TestMsgDeleteMarkerMaxAge(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name: "foo", Subjects: []string{"FOO.*"}, AllowMsgTTL: true, SubjectDeleteMarkerTTL: 50 * time.Second, MaxAge: 1 * time.Second})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	_, err = js.Publish(ctx, "FOO.1", []byte("msg1"))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	time.Sleep(1500 * time.Millisecond)
	gotMsg, err := stream.GetLastMsgForSubject(ctx, "FOO.1")
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if ttlMarker := gotMsg.Header.Get("Nats-Marker-Reason"); ttlMarker != "MaxAge" {
		t.Fatalf("Expected message to have Marker-Reason header set to MaxAge; got: %s", ttlMarker)
	}
	if ttl := gotMsg.Header.Get("Nats-TTL"); ttl != "50s" {
		t.Fatalf("Expected message to have Nats-TTL header set to 50s; got: %s", ttl)
	}
}

func TestPublishAsyncWithTTL(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name: "foo", Subjects: []string{"FOO.*"}, MaxMsgSize: 64, AllowMsgTTL: true})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	paf, err := js.PublishAsync("FOO.1", []byte("msg"), jetstream.WithMsgTTL(1*time.Second))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	var ack *jetstream.PubAck
	select {
	case ack = <-paf.Ok():
	case <-time.After(5 * time.Second):
		t.Fatalf("Did not receive ack")
	}

	gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if ttl := gotMsg.Header.Get("Nats-TTL"); ttl != "1s" {
		t.Fatalf("Expected message to have TTL header set to 1s; got: %s", ttl)
	}
	time.Sleep(1500 * time.Millisecond)
	_, err = stream.GetMsg(ctx, ack.Sequence)
	if !errors.Is(err, jetstream.ErrMsgNotFound) {
		t.Fatalf("Expected not found error; got: %v", err)
	}
}

func TestPublish(t *testing.T) {
	// Only very basic test cases, as most use cases are tested in TestPublishMsg
	tests := []struct {
		name      string
		msg       []byte
		subject   string
		opts      []jetstream.PublishOpt
		withError error
	}{
		{
			name:    "publish single message on stream, no options",
			msg:     []byte("msg"),
			subject: "FOO.1",
		},
		{
			name:    "publish single message on stream with message id",
			msg:     []byte("msg"),
			subject: "FOO.1",
			opts:    []jetstream.PublishOpt{jetstream.WithMsgID("1")},
		},
		{
			name:      "empty subject passed",
			msg:       []byte("msg"),
			subject:   "",
			withError: nats.ErrBadSubject,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			srv := RunBasicJetStreamServer()
			defer shutdownJSServerAndRemoveStorage(t, srv)
			nc, err := nats.Connect(srv.ClientURL())
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}

			js, err := jetstream.New(nc)
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			defer nc.Close()

			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}, MaxMsgSize: 64})
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}

			ack, err := js.Publish(ctx, test.subject, test.msg, test.opts...)
			if test.withError != nil {
				if err == nil || !errors.Is(err, test.withError) {
					t.Fatalf("Expected error: %v; got: %v", test.withError, err)
				}
				return
			}
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			if ack.Sequence != 1 || ack.Stream != "foo" {
				t.Fatalf("Invalid ack; want sequence 1 on stream foo, got: %v", ack)
			}
		})
	}
}

func TestPublishTimeout(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	js, err := jetstream.New(nc, jetstream.WithDefaultTimeout(200*time.Millisecond))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	// create stream with no ack to force timeout
	_, err = js.CreateStream(context.Background(), jetstream.StreamConfig{
		Name:     "foo",
		Subjects: []string{"FOO.*"},
		NoAck:    true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	now := time.Now()
	_, err = js.Publish(context.Background(), "FOO.1", []byte("msg"))
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("Expected deadline exceeded error; got: %v", err)
	}
	since := time.Since(now)
	if since < 200*time.Millisecond || since > 500*time.Millisecond {
		t.Fatalf("Expected timeout to be around 200ms; got: %v", since)
	}
}

func TestPublishMsgAsync(t *testing.T) {
	type publishConfig struct {
		msg              *nats.Msg
		opts             []jetstream.PublishOpt
		expectedHeaders  nats.Header
		expectedAck      jetstream.PubAck
		withAckError     func(*testing.T, error)
		withPublishError func(*testing.T, error)
	}
	tests := []struct {
		name      string
		msgs      []publishConfig
		srvConfig []byte
		timeout   time.Duration
	}{
		{
			name: "publish 3 simple messages, no opts",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
						Domain:   "",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
						Domain:   "",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
						Domain:   "",
					},
				},
			},
		},
		{
			name: "publish with ack timeout set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
						Domain:   "",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
						Domain:   "",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
						Domain:   "",
					},
				},
			},
			timeout: time.Second,
		},
		{
			name: "publish 3 messages with message ID, with duplicate",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:    "foo",
						Sequence:  1,
						Duplicate: true,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"2"},
					},
				},
			},
		},
		{
			name: "expect last msg ID",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"2"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastMsgID("2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Msg-Id": []string{"2"},
					},
				},
			},
		},
		{
			name: "invalid last msg id",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithMsgID("1")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Msg-Id": []string{"1"},
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastMsgID("abc")},
					withAckError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10070 {
							t.Fatalf("Expected error code: 10070; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect last sequence",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequence(2)},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Sequence": []string{"2"},
					},
				},
			},
		},
		{
			name: "invalid last seq",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequence(123)},
					withAckError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10071 {
							t.Fatalf("Expected error code: 10071; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect last sequence per subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequencePerSubject(1)},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Subject-Sequence": []string{"1"},
					},
				},
			},
		},
		{
			name: "invalid last sequence per subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequencePerSubject(123)},
					withAckError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10071 {
							t.Fatalf("Expected error code: 10071; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect last sequence for subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequenceForSubject(2, "FOO.2")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Last-Subject-Sequence":         []string{"2"},
						"Nats-Expected-Last-Subject-Sequence-Subject": []string{"FOO.2"},
					},
				},
			},
		},
		{
			name: "invalid last sequence for subject",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.2",
					},
					opts: []jetstream.PublishOpt{},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectLastSequenceForSubject(123, "FOO.2")},
					withAckError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10071 {
							t.Fatalf("Expected error code: 10071; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "expect stream header",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectStream("foo")},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
					},
					expectedHeaders: nats.Header{
						"Nats-Expected-Stream": []string{"foo"},
					},
				},
			},
		},
		{
			name: "invalid expected stream header",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithExpectStream("abc")},
					withAckError: func(t *testing.T, err error) {
						var apiErr *jetstream.APIError
						if ok := errors.As(err, &apiErr); !ok {
							t.Fatalf("Expected API error; got: %v", err)
						}
						if apiErr.ErrorCode != 10060 {
							t.Fatalf("Expected error code: 10060; got: %d", apiErr.ErrorCode)
						}
					},
				},
			},
		},
		{
			name: "publish 3 simple messages with domain set",
			srvConfig: []byte(
				`
					listen: 127.0.0.1:-1
					jetstream: {domain: "test-domain"}
				`),
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 1,
						Domain:   "test-domain",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 2"),
						Subject: "FOO.1",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 2,
						Domain:   "test-domain",
					},
				},
				{
					msg: &nats.Msg{
						Data:    []byte("msg 3"),
						Subject: "FOO.2",
					},
					expectedAck: jetstream.PubAck{
						Stream:   "foo",
						Sequence: 3,
						Domain:   "test-domain",
					},
				},
			},
		},
		{
			name: "invalid subject set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "ABC",
					},
					withAckError: func(t *testing.T, err error) {
						if !errors.Is(err, jetstream.ErrNoStreamResponse) {
							t.Fatalf("Expected error: %v; got: %v", jetstream.ErrNoStreamResponse, err)
						}
					},
				},
			},
		},
		{
			name: "invalid retry number set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithRetryAttempts(-1)},
					withPublishError: func(t *testing.T, err error) {
						if !errors.Is(err, jetstream.ErrInvalidOption) {
							t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
						}
					},
				},
			},
		},
		{
			name: "invalid retry wait set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithRetryWait(-1)},
					withPublishError: func(t *testing.T, err error) {
						if !errors.Is(err, jetstream.ErrInvalidOption) {
							t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
						}
					},
				},
			},
		},
		{
			name: "invalid stall wait set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
					},
					opts: []jetstream.PublishOpt{jetstream.WithStallWait(-1)},
					withPublishError: func(t *testing.T, err error) {
						if !errors.Is(err, jetstream.ErrInvalidOption) {
							t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
						}
					},
				},
			},
		},
		{
			name: "reply subject set",
			msgs: []publishConfig{
				{
					msg: &nats.Msg{
						Data:    []byte("msg 1"),
						Subject: "FOO.1",
						Reply:   "BAR",
					},
					withPublishError: func(t *testing.T, err error) {
						if !errors.Is(err, jetstream.ErrAsyncPublishReplySubjectSet) {
							t.Fatalf("Expected error: %v; got: %v", nats.ErrNoResponders, err)
						}
					},
				},
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var srv *server.Server
			if test.srvConfig != nil {
				conf := createConfFile(t, test.srvConfig)
				defer os.Remove(conf)
				srv, _ = RunServerWithConfig(conf)
			} else {
				srv = RunBasicJetStreamServer()
			}
			defer shutdownJSServerAndRemoveStorage(t, srv)
			nc, err := nats.Connect(srv.ClientURL())
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}

			opts := []jetstream.JetStreamOpt{}
			if test.timeout != 0 {
				opts = append(opts, jetstream.WithPublishAsyncTimeout(test.timeout))
			}
			js, err := jetstream.New(nc, opts...)
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			defer nc.Close()

			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}, MaxMsgSize: 128})
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}

			for _, pub := range test.msgs {
				ackFuture, err := js.PublishMsgAsync(pub.msg, pub.opts...)
				if pub.withPublishError != nil {
					pub.withPublishError(t, err)
					continue
				}
				select {
				case ack := <-ackFuture.Ok():
					if pub.withAckError != nil {
						t.Fatalf("Expected error, got nil")
					}
					if *ack != pub.expectedAck {
						t.Fatalf("Invalid ack received; want: %v; got: %v", pub.expectedAck, ack)
					}
					msg := ackFuture.Msg()
					if !reflect.DeepEqual(pub.expectedHeaders, msg.Header) {
						t.Fatalf("Invalid headers on message; want: %v; got: %v", pub.expectedHeaders, pub.msg.Header)
					}
					if string(msg.Data) != string(pub.msg.Data) {
						t.Fatalf("Invalid message in ack; want: %q; got: %q", string(pub.msg.Data), string(msg.Data))
					}
				case err := <-ackFuture.Err():
					if pub.withAckError == nil {
						t.Fatalf("Expected no error. got: %v", err)
					}
					pub.withAckError(t, err)
				}
			}

			select {
			case <-js.PublishAsyncComplete():
			case <-time.After(5 * time.Second):
				t.Fatalf("Did not receive completion signal")
			}
		})
	}
}

func TestPublishMsgAsyncWithPendingMsgs(t *testing.T) {
	t.Run("outstanding ack exceed limit", func(t *testing.T) {
		srv := RunBasicJetStreamServer()
		defer shutdownJSServerAndRemoveStorage(t, srv)
		nc, err := nats.Connect(srv.ClientURL())
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		js, err := jetstream.New(nc, jetstream.WithPublishAsyncMaxPending(5))
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		defer nc.Close()

		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		for range 20 {
			_, err = js.PublishAsync("FOO.1", []byte("msg"))
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			if numPending := js.PublishAsyncPending(); numPending > 5 {
				t.Fatalf("Expected 5 pending messages, got: %d", numPending)
			}
		}
	})
	t.Run("too many messages without ack", func(t *testing.T) {
		srv := RunBasicJetStreamServer()
		defer shutdownJSServerAndRemoveStorage(t, srv)
		nc, err := nats.Connect(srv.ClientURL())
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		js, err := jetstream.New(nc, jetstream.WithPublishAsyncMaxPending(5))
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		defer nc.Close()

		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		_, err = js.CreateStream(ctx, jetstream.StreamConfig{
			Name:     "foo",
			Subjects: []string{"FOO.*"},
			// disable stream acks
			NoAck: true,
		})
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		for range 5 {
			_, err = js.PublishAsync("FOO.1", []byte("msg"), jetstream.WithStallWait(10*time.Millisecond))
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
		}
		if _, err = js.PublishAsync("FOO.1", []byte("msg"), jetstream.WithStallWait(10*time.Millisecond)); err == nil || !errors.Is(err, jetstream.ErrTooManyStalledMsgs) {
			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrTooManyStalledMsgs, err)
		}
	})

	t.Run("with server restart", func(t *testing.T) {
		srv := RunBasicJetStreamServer()
		defer shutdownJSServerAndRemoveStorage(t, srv)
		nc, err := nats.Connect(srv.ClientURL())
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		js, err := jetstream.New(nc)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		defer nc.Close()

		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		go func() {
			for range 50 {
				_, _ = js.PublishAsync("FOO.1", []byte("msg"))
			}
		}()
		srv = restartBasicJSServer(t, srv)
		defer shutdownJSServerAndRemoveStorage(t, srv)

		select {
		case <-js.PublishAsyncComplete():
		case <-time.After(10 * time.Second):
			t.Fatalf("Did not receive completion signal")
		}
	})
}

func TestPublishAsyncResetPendingOnReconnect(t *testing.T) {
	s := RunBasicJetStreamServer()

	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	errs := make(chan error, 1)
	done := make(chan struct{}, 1)
	acks := make(chan jetstream.PubAckFuture, 100)
	wg := sync.WaitGroup{}
	go func() {
		for range 100 {
			if ack, err := js.PublishAsync("FOO.A", []byte("hello")); err != nil {
				errs <- err
				return
			} else {
				acks <- ack
			}
			wg.Add(1)
		}
		close(acks)
		done <- struct{}{}
	}()
	select {
	case <-done:
	case err := <-errs:
		t.Fatalf("Unexpected error during publish: %v", err)
	case <-time.After(5 * time.Second):
		t.Fatalf("Did not receive completion signal")
	}
	for ack := range acks {
		go func(paf jetstream.PubAckFuture) {
			select {
			case <-paf.Ok():
			case err := <-paf.Err():
				if !errors.Is(err, nats.ErrDisconnected) && !errors.Is(err, nats.ErrNoResponders) {
					errs <- fmt.Errorf("Expected error: %v or %v; got: %v", nats.ErrDisconnected, nats.ErrNoResponders, err)
				}
			case <-time.After(5 * time.Second):
				errs <- errors.New("Did not receive completion signal")
			}
			wg.Done()
		}(ack)
	}
	s = restartBasicJSServer(t, s)
	defer shutdownJSServerAndRemoveStorage(t, s)

	wg.Wait()
	select {
	case err := <-errs:
		t.Fatalf("Unexpected error: %v", err)
	default:
	}
}

func TestPublishAsyncRetry(t *testing.T) {
	tests := []struct {
		name     string
		pubOpts  []jetstream.PublishOpt
		ackError error
	}{
		{
			name: "retry until stream is ready",
			pubOpts: []jetstream.PublishOpt{
				jetstream.WithRetryAttempts(10),
				jetstream.WithRetryWait(100 * time.Millisecond),
			},
		},
		{
			name: "fail after max retries",
			pubOpts: []jetstream.PublishOpt{
				jetstream.WithRetryAttempts(2),
				jetstream.WithRetryWait(50 * time.Millisecond),
			},
			ackError: jetstream.ErrNoStreamResponse,
		},
		{
			name: "retries disabled",
			pubOpts: []jetstream.PublishOpt{
				jetstream.WithRetryAttempts(0),
			},
			ackError: jetstream.ErrNoStreamResponse,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			s := RunBasicJetStreamServer()
			defer shutdownJSServerAndRemoveStorage(t, s)

			nc, err := nats.Connect(s.ClientURL())
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}

			// set max pending to 1 so that we can test if retries don't cause stall
			js, err := jetstream.New(nc, jetstream.WithPublishAsyncMaxPending(1))
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			defer nc.Close()
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()

			test.pubOpts = append(test.pubOpts, jetstream.WithStallWait(1*time.Nanosecond))
			ack, err := js.PublishAsync("foo", []byte("hello"), test.pubOpts...)
			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			publishComplete := js.PublishAsyncComplete()
			errs := make(chan error, 1)
			go func() {
				// create stream with delay so that publish will receive no responders
				time.Sleep(300 * time.Millisecond)
				if _, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "TEST", Subjects: []string{"foo"}}); err != nil {
					errs <- err
				}
			}()
			select {
			case <-ack.Ok():
			case err := <-ack.Err():
				if test.ackError != nil {
					if !errors.Is(err, test.ackError) {
						t.Fatalf("Expected error: %v; got: %v", test.ackError, err)
					}
				} else {
					t.Fatalf("Unexpected ack error: %v", err)
				}
			case err := <-errs:
				t.Fatalf("Error creating stream: %v", err)
			case <-time.After(5 * time.Second):
				t.Fatalf("Timeout waiting for ack")
			}

			select {
			case <-publishComplete:
			case <-time.After(5 * time.Second):
				t.Fatalf("Did not receive completion signal")
			}
		})
	}
}

func TestPublishAsyncRetryInErrHandler(t *testing.T) {
	s := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, s)

	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	streamCreated := make(chan struct{})
	errCB := func(js jetstream.JetStream, m *nats.Msg, e error) {
		<-streamCreated
		_, err := js.PublishMsgAsync(m)
		if err != nil {
			t.Fatalf("Unexpected error when republishing: %v", err)
		}
	}

	js, err := jetstream.New(nc, jetstream.WithPublishAsyncErrHandler(errCB))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	errs := make(chan error, 1)
	done := make(chan struct{}, 1)
	go func() {
		for range 10 {
			if _, err := js.PublishAsync("FOO.A", []byte("hello"), jetstream.WithRetryAttempts(0)); err != nil {
				errs <- err
				return
			}
		}
		done <- struct{}{}
	}()
	select {
	case <-done:
	case err := <-errs:
		t.Fatalf("Unexpected error during publish: %v", err)
	case <-time.After(5 * time.Second):
		t.Fatalf("Did not receive completion signal")
	}
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	close(streamCreated)
	select {
	case <-js.PublishAsyncComplete():
	case <-time.After(5 * time.Second):
		t.Fatalf("Did not receive completion signal")
	}

	info, err := stream.Info(context.Background())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	if info.State.Msgs != 10 {
		t.Fatalf("Expected 10 messages in the stream; got: %d", info.State.Msgs)
	}
}

func TestPublishAsyncAckTimeout(t *testing.T) {
	s := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, s)

	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	errs := make(chan error, 1)
	js, err := jetstream.New(nc,
		jetstream.WithPublishAsyncTimeout(50*time.Millisecond),
		jetstream.WithPublishAsyncErrHandler(func(js jetstream.JetStream, m *nats.Msg, e error) {
			errs <- e
		}),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	_, err = js.CreateStream(context.Background(), jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}, NoAck: true})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ack, err := js.PublishAsync("FOO.A", []byte("hello"))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	select {
	case <-ack.Ok():
		t.Fatalf("Expected timeout")
	case err := <-ack.Err():
		if !errors.Is(err, jetstream.ErrAsyncPublishTimeout) {
			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrAsyncPublishTimeout, err)
		}
	case <-time.After(time.Second):
		t.Fatalf("Did not receive ack timeout")
	}

	// check if error callback is called
	select {
	case err := <-errs:
		if !errors.Is(err, jetstream.ErrAsyncPublishTimeout) {
			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrAsyncPublishTimeout, err)
		}
	case <-time.After(time.Second):
		t.Fatalf("Did not receive error from error handler")
	}

	if js.PublishAsyncPending() != 0 {
		t.Fatalf("Expected no pending messages")
	}

	select {
	case <-js.PublishAsyncComplete():
	case <-time.After(100 * time.Millisecond):
		t.Fatalf("Did not receive completion signal")
	}
}

func TestPublishAsyncAckHandler(t *testing.T) {
	s := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, s)

	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	acks := make(chan *jetstream.PubAck, 1)
	js, err := jetstream.New(nc, jetstream.WithPublishAsyncAckHandler(func(_ jetstream.JetStream, _ *nats.Msg, ack *jetstream.PubAck) {
		acks <- ack
	}))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	cfg := jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}
	if _, err := js.CreateStream(t.Context(), cfg); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	paf, err := js.PublishAsync("FOO.A", []byte("flightless, not fightless"))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	select {
	case <-time.After(time.Second):
		t.Fatalf("Did not receive ack from ack handler")
	case <-acks: // the ack handler must fire when the publish is successful
	}

	select {
	case <-time.After(time.Second):
		t.Fatalf("Did not receive ack from PubAckFuture")
	case err := <-paf.Err():
		t.Fatalf("Unexpected error: %v", err)
	case <-paf.Ok(): // PubAckFuture must still resolve independently of the ack handler
	}
}

func TestPublishAsyncAckHandlerSkippedOnTimeout(t *testing.T) {
	s := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, s)

	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	acks, errs := make(chan *jetstream.PubAck, 1), make(chan error, 1)
	js, err := jetstream.New(nc,
		jetstream.WithPublishAsyncTimeout(50*time.Millisecond),
		jetstream.WithPublishAsyncAckHandler(func(_ jetstream.JetStream, _ *nats.Msg, ack *jetstream.PubAck) {
			acks <- ack
		}),
		jetstream.WithPublishAsyncErrHandler(func(_ jetstream.JetStream, _ *nats.Msg, err error) {
			errs <- err
		}),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	cfg := jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}, NoAck: true}
	if _, err := js.CreateStream(t.Context(), cfg); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if _, err := js.PublishAsync("FOO.A", []byte("flightless, not fightless")); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	select {
	case err := <-errs:
		if !errors.Is(err, jetstream.ErrAsyncPublishTimeout) {
			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrAsyncPublishTimeout, err)
		}
	case <-time.After(time.Second):
		t.Fatalf("Did not receive error from error handler")
	}

	select {
	case ack := <-acks:
		t.Fatalf("Did not expect ack handler to be called; got: %v", ack)
	case <-time.After(100 * time.Millisecond): // the ack handler must not fire when the publish errors out.
	}
}

func TestPublishAsyncClearStall(t *testing.T) {
	s := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, s)

	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	js, err := jetstream.New(nc,
		jetstream.WithPublishAsyncTimeout(500*time.Millisecond),
		jetstream.WithPublishAsyncMaxPending(100))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	// use stream with no acks to test stalling
	_, err = js.CreateStream(context.Background(), jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}, NoAck: true})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	for range 100 {
		_, err := js.PublishAsync("FOO.A", []byte("hello"), jetstream.WithStallWait(1*time.Nanosecond))
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
	}
	// after publishing 100 messages, next one should fail with ErrTooManyStalledMsgs
	_, err = js.PublishAsync("FOO.A", []byte("hello"), jetstream.WithStallWait(50*time.Millisecond))
	if !errors.Is(err, jetstream.ErrTooManyStalledMsgs) {
		t.Fatalf("Expected error: %v; got: %v", jetstream.ErrTooManyStalledMsgs, err)
	}

	// after publish timeout all pending messages should be cleared
	// and we should be able to publish again
	select {
	case <-js.PublishAsyncComplete():
	case <-time.After(2 * time.Second):
		t.Fatalf("Did not receive completion signal")
	}

	if _, err = js.PublishAsync("FOO.A", []byte("hello")); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if js.PublishAsyncPending() != 1 {
		t.Fatalf("Expected 1 pending message; got: %d", js.PublishAsyncPending())
	}
}

func TestPublishWithScheduleAt(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	cons, err := stream.CreateConsumer(ctx, jetstream.ConsumerConfig{
		FilterSubject: "target.>",
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	schedTime := time.Now().Add(time.Second).UTC().Truncate(time.Second)
	ack, err := js.Publish(ctx, "schedule.at", []byte("hello"),
		jetstream.WithScheduleAt(schedTime),
		jetstream.WithScheduleTarget("target.at"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleHeader); got != "@at "+schedTime.Format(time.RFC3339) {
		t.Fatalf("Expected schedule header %q; got: %q", "@at "+schedTime.Format(time.RFC3339), got)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleTargetHeader); got != "target.at" {
		t.Fatalf("Expected schedule target header %q; got: %q", "target.at", got)
	}

	msg, err := cons.Next(jetstream.FetchMaxWait(5 * time.Second))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if string(msg.Data()) != "hello" {
		t.Fatalf("Expected message data %q; got: %q", "hello", string(msg.Data()))
	}
}

func TestPublishWithScheduleEvery(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ack, err := js.Publish(ctx, "schedule.every", nil,
		jetstream.WithScheduleEvery(5*time.Second),
		jetstream.WithScheduleTarget("target.every"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	// check if correct headers are set
	gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleHeader); got != "@every 5s" {
		t.Fatalf("Expected schedule header %q; got: %q", "@every 5s", got)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleTargetHeader); got != "target.every" {
		t.Fatalf("Expected schedule target header %q; got: %q", "target.every", got)
	}

	// invalid duration (should be at least 1s)
	_, err = js.Publish(ctx, "schedule.every", nil,
		jetstream.WithScheduleEvery(500*time.Millisecond),
		jetstream.WithScheduleTarget("target.every"),
	)
	if !errors.Is(err, jetstream.ErrInvalidOption) {
		t.Fatalf("Expected ErrInvalidOption; got: %v", err)
	}
}

func TestPublishWithScheduleMissingTarget(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err = js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	_, err = js.Publish(ctx, "schedule.notarget", nil,
		jetstream.WithScheduleAt(time.Now().Add(1*time.Hour)),
	)
	if !errors.Is(err, jetstream.ErrScheduleTargetInvalid) {
		t.Fatalf("Expected ErrScheduleTargetInvalid; got: %v", err)
	}
}

func TestPublishWithScheduleCron(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	t.Run("6-field cron", func(t *testing.T) {
		ack, err := js.Publish(ctx, "schedule.cron1", nil,
			jetstream.WithScheduleCron("0 30 * * * *"),
			jetstream.WithScheduleTarget("target.cron1"),
		)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		// check if correct headers are set
		gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		if got := gotMsg.Header.Get(jetstream.ScheduleHeader); got != "0 30 * * * *" {
			t.Fatalf("Expected schedule header %q; got: %q", "0 30 * * * *", got)
		}
	})

	t.Run("predefined schedule", func(t *testing.T) {
		ack, err := js.Publish(ctx, "schedule.cron2", nil,
			jetstream.WithScheduleCron(jetstream.ScheduleHourly),
			jetstream.WithScheduleTarget("target.cron2"),
		)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}

		// check if correct headers are set
		gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		if got := gotMsg.Header.Get(jetstream.ScheduleHeader); got != "@hourly" {
			t.Fatalf("Expected schedule header %q; got: %q", "@hourly", got)
		}
	})
}

func TestPublishWithScheduleSource(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>", "source.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	cons, err := js.CreateConsumer(ctx, "SCHED", jetstream.ConsumerConfig{
		FilterSubject: "target.>",
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ack, err := js.Publish(ctx, "schedule.sample", nil,
		jetstream.WithScheduleEvery(time.Second),
		jetstream.WithScheduleTarget("target.sample"),
		jetstream.WithScheduleSource("source.data"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleSourceHeader); got != "source.data" {
		t.Fatalf("Expected schedule source header %q; got: %q", "source.data", got)
	}

	if _, err := js.Publish(ctx, "source.data", []byte("hello")); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	msg, err := cons.Next(jetstream.FetchMaxWait(5 * time.Second))
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if string(msg.Data()) != "hello" {
		t.Fatalf("Expected message data %q; got: %q", "hello", string(msg.Data()))
	}

}

func TestPublishWithScheduleTTL(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
		AllowMsgTTL:       true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	t.Run("duration", func(t *testing.T) {
		ack, err := js.Publish(ctx, "schedule.ttl1", nil,
			jetstream.WithScheduleEvery(5*time.Second),
			jetstream.WithScheduleTarget("target.ttl1"),
			jetstream.WithScheduleTTL(10*time.Minute),
		)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		if got := gotMsg.Header.Get(jetstream.ScheduleTTLHeader); got != "10m0s" {
			t.Fatalf("Expected schedule TTL header %q; got: %q", "10m0s", got)
		}
	})

	t.Run("never", func(t *testing.T) {
		ack, err := js.Publish(ctx, "schedule.ttl2", nil,
			jetstream.WithScheduleEvery(5*time.Second),
			jetstream.WithScheduleTarget("target.ttl2"),
			jetstream.WithScheduleTTLNever(),
		)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
		if err != nil {
			t.Fatalf("Unexpected error: %v", err)
		}
		if got := gotMsg.Header.Get(jetstream.ScheduleTTLHeader); got != "never" {
			t.Fatalf("Expected schedule TTL header %q; got: %q", "never", got)
		}
	})
}

func TestPublishWithScheduleTimeZone(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ack, err := js.Publish(ctx, "schedule.tz", nil,
		jetstream.WithScheduleCron("0 0 9 * * *"),
		jetstream.WithScheduleTarget("target.tz"),
		jetstream.WithScheduleTimeZone("America/New_York"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleTimeZoneHeader); got != "America/New_York" {
		t.Fatalf("Expected schedule time zone header %q; got: %q", "America/New_York", got)
	}
}

func TestPublishAsyncWithSchedule(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	paf, err := js.PublishAsync("schedule.async", nil,
		jetstream.WithScheduleEvery(5*time.Second),
		jetstream.WithScheduleTarget("target.async"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	var ack *jetstream.PubAck
	select {
	case ack = <-paf.Ok():
	case <-time.After(5 * time.Second):
		t.Fatalf("Did not receive ack")
	}

	gotMsg, err := stream.GetMsg(ctx, ack.Sequence)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleHeader); got != "@every 5s" {
		t.Fatalf("Expected schedule header %q; got: %q", "@every 5s", got)
	}
	if got := gotMsg.Header.Get(jetstream.ScheduleTargetHeader); got != "target.async" {
		t.Fatalf("Expected schedule target header %q; got: %q", "target.async", got)
	}
}

func TestScheduleAtDelivery(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	_, err = js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	cons, err := js.CreateConsumer(ctx, "SCHED", jetstream.ConsumerConfig{
		FilterSubject: "target.at",
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	_, err = js.Publish(ctx, "schedule.at", []byte("scheduled payload"),
		jetstream.WithScheduleAt(time.Now().Add(2*time.Second)),
		jetstream.WithScheduleTarget("target.at"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	msg, err := cons.Next(jetstream.FetchMaxWait(5 * time.Second))
	if err != nil {
		t.Fatalf("Expected to receive scheduled message: %v", err)
	}
	if string(msg.Data()) != "scheduled payload" {
		t.Fatalf("Expected payload %q; got: %q", "scheduled payload", string(msg.Data()))
	}
	if msg.Headers().Get(jetstream.SchedulerHeader) == "" {
		t.Fatal("Expected Nats-Scheduler header to be set on delivered message")
	}
}

func TestScheduleEveryDelivery(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	_, err = js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	cons, err := js.CreateConsumer(ctx, "SCHED", jetstream.ConsumerConfig{
		FilterSubject: "target.every",
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	_, err = js.Publish(ctx, "schedule.every", []byte("repeating"),
		jetstream.WithScheduleEvery(1*time.Second),
		jetstream.WithScheduleTarget("target.every"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	// Verify at least 2 deliveries to confirm repeating behavior.
	for range 2 {
		msg, err := cons.Next(jetstream.FetchMaxWait(5 * time.Second))
		if err != nil {
			t.Fatalf("Expected to receive scheduled message: %v", err)
		}
		if string(msg.Data()) != "repeating" {
			t.Fatalf("Expected payload %q; got: %q", "repeating", string(msg.Data()))
		}
	}
}

func TestScheduleCronDelivery(t *testing.T) {
	srv := RunBasicJetStreamServer()
	defer shutdownJSServerAndRemoveStorage(t, srv)
	nc, err := nats.Connect(srv.ClientURL())
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	_, err = js.CreateStream(ctx, jetstream.StreamConfig{
		Name:              "SCHED",
		Subjects:          []string{"schedule.>", "target.>"},
		AllowMsgSchedules: true,
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	cons, err := js.CreateConsumer(ctx, "SCHED", jetstream.ConsumerConfig{
		FilterSubject: "target.cron",
	})
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	// Schedule every second using 6-field cron.
	_, err = js.Publish(ctx, "schedule.cron", []byte("cron payload"),
		jetstream.WithScheduleCron("* * * * * *"),
		jetstream.WithScheduleTarget("target.cron"),
	)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	msg, err := cons.Next(jetstream.FetchMaxWait(5 * time.Second))
	if err != nil {
		t.Fatalf("Expected to receive scheduled message: %v", err)
	}
	if string(msg.Data()) != "cron payload" {
		t.Fatalf("Expected payload %q; got: %q", "cron payload", string(msg.Data()))
	}
	if next := msg.Headers().Get(jetstream.ScheduleNextHeader); next == "" {
		t.Fatal("Expected Nats-Schedule-Next header to be set on delivered message")
	}
}