fix atomic batch publish test

Signed-off-by: Piotr Piotrowski <piotr@synadia.com>

Author: piotrpio

jetstreamext/errors.go

@@ -22,6 +22,10 @@ const (
 	JSErrCodeFastBatchInvalidPattern jetstream.ErrorCode = 10204
 	JSErrCodeFastBatchInvalidID      jetstream.ErrorCode = 10205
 	JSErrCodeFastBatchUnknownID      jetstream.ErrorCode = 10206
+
+	// Too many inflight error codes
+	JSErrCodeAtomicPublishTooManyInflight jetstream.ErrorCode = 10210
+	JSErrCodeBatchPublishTooManyInflight  jetstream.ErrorCode = 10211
 )
 
 var (
@@ -51,6 +55,12 @@ var (
 	// ErrBatchPublishInvalidGapMode is returned when invalid batch gap mode is specified.
 	ErrBatchPublishInvalidGapMode jetstream.JetStreamError = &jsError{apiErr: &jetstream.APIError{ErrorCode: JSErrCodeBatchPublishInvalidGapMode, Description: "invalid batch gap mode", Code: 400}}
 
+	// ErrAtomicPublishTooManyInflight is returned when there are too many inflight atomic publish batches.
+	ErrAtomicPublishTooManyInflight jetstream.JetStreamError = &jsError{apiErr: &jetstream.APIError{ErrorCode: JSErrCodeAtomicPublishTooManyInflight, Description: "atomic publish too many inflight", Code: 429}}
+
+	// ErrBatchPublishTooManyInflight is returned when there are too many inflight batch publish batches.
+	ErrBatchPublishTooManyInflight jetstream.JetStreamError = &jsError{apiErr: &jetstream.APIError{ErrorCode: JSErrCodeBatchPublishTooManyInflight, Description: "batch publish too many inflight", Code: 429}}
+
 	// ErrBatchClosed is returned when attempting to use a batch that has been closed.
 	ErrBatchClosed = &jsError{message: "batch publisher closed"}
 

jetstreamext/test/publishbatch_test.go

@@ -310,16 +310,16 @@ func TestBatchPublisher(t *testing.T) {
 		err = batch.Add("test.1", []byte("message 1"))
 		if err != nil {
 			// With flow control enabled (WaitFirst=true by default), the first Add may fail
-			if errors.Is(err, jetstreamext.ErrBatchPublishIncomplete) {
+			if errors.Is(err, jetstreamext.ErrBatchPublishIncomplete) || errors.Is(err, jetstreamext.ErrAtomicPublishTooManyInflight) || errors.Is(err, jetstreamext.ErrBatchPublishTooManyInflight) {
 				// This is expected - too many outstanding batches
 				return
 			}
 			t.Fatalf("Unexpected error adding message to batch: %v", err)
 		}
 		// If Add didn't fail, Commit should fail
 		_, err = batch.Commit(ctx, "test.2", []byte("message 2"))
-		if !errors.Is(err, jetstreamext.ErrBatchPublishIncomplete) {
-			t.Fatalf("Expected ErrBatchPublishIncomplete when too many outstanding batches, got %v", err)
+		if !errors.Is(err, jetstreamext.ErrBatchPublishIncomplete) && !errors.Is(err, jetstreamext.ErrAtomicPublishTooManyInflight) && !errors.Is(err, jetstreamext.ErrBatchPublishTooManyInflight) {
+			t.Fatalf("Expected too many inflight error when too many outstanding batches, got %v", err)
 		}
 	})