Repository: numaproj/numaflow
Pull request #418 (2022-12-09 01:52:58)

refactor: nats/jetstream testing
Signed-off-by: Derek Wang <whynowy@gmail.com>

Explain what this PR does.

<!--

Before you push your changes:

* Run `make pre-push -B` to fix codegen and lint problems.

Then, you MUST:

* Sign-off your commit (otherwise the DCO check will fail).
* Use [a conventional commit message](https://www.conventionalcommits.org/en/v1.0.0/) (otherwise the commit message check will fail).

If you did not do this, reset all your commit and replace them with a single commit:

```
git reset HEAD~1 ;# change 1 to how many commits you made
git commit --signoff -m 'feat: my feat. Fixes #1234'
```

When creating your PR: 

* Make sure that "Fixes #" or "Closes #" is in both the PR title (for release notes) and description (to automatically link and close the issue).
* Say how you tested your changes. If you changed the UI, attach screenshots.
* Set your PR as a draft initially.
* Your PR needs to pass the required checks before it can be approved. 
* Once required tests have passed, mark your PR "Ready for review".

If changes were requested, once you've made them, you MUST dismiss the review to get it reviewed again.

-->



Patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 00a91e5f..58616c1c 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -107,13 +107,6 @@ jobs:
           key: ${{ runner.os }}-go-build-v1-${{ github.run_id }}
       - name: Get dependencies
         run: go mod download
-      - name: Install JetStream
-        run: |
-          set -eux -o pipefail
-          JS_VERSION=2.9.6
-          curl -L https://github.com/nats-io/nats-server/releases/download/v$JS_VERSION/nats-server-v$JS_VERSION-linux-amd64.zip -o nats-server.zip
-          sudo unzip nats-server.zip -d nats-server
-          nats-server/nats-server-v$JS_VERSION-linux-amd64/nats-server -js &
       - name: Test
         run: make test-code
 
diff --git a/Makefile b/Makefile
index 90f4fcc4..b5ec399a 100644
--- a/Makefile
+++ b/Makefile
@@ -90,14 +90,14 @@ test-coverage:
 
 .PHONY: test-coverage-with-isb
 test-coverage-with-isb:
-	go test -covermode=atomic -coverprofile=test/profile.cov.tmp -tags=isb_redis,isb_jetstream $(shell go list ./... | grep -v /vendor/ | grep -v /numaflow/test/ | grep -v /pkg/client/ | grep -v /pkg/proto/ | grep -v /hack/)
+	go test -covermode=atomic -coverprofile=test/profile.cov.tmp -tags=isb_redis $(shell go list ./... | grep -v /vendor/ | grep -v /numaflow/test/ | grep -v /pkg/client/ | grep -v /pkg/proto/ | grep -v /hack/)
 	cat test/profile.cov.tmp | grep -v v1alpha1/zz_generated | grep -v v1alpha1/generated > test/profile.cov
 	rm test/profile.cov.tmp
 	go tool cover -func=test/profile.cov
 
 .PHONY: test-code
 test-code:
-	go test -tags=isb_redis,isb_jetstream -race -v $(shell go list ./... | grep -v /vendor/ | grep -v /numaflow/test/)
+	go test -tags=isb_redis -race -v $(shell go list ./... | grep -v /vendor/ | grep -v /numaflow/test/)
 
 test-e2e:
 test-kafka-e2e:
diff --git a/pkg/daemon/server/daemon_server.go b/pkg/daemon/server/daemon_server.go
index 1d1b9987..978d6f9a 100644
--- a/pkg/daemon/server/daemon_server.go
+++ b/pkg/daemon/server/daemon_server.go
@@ -36,7 +36,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/apis/proto/daemon"
 	"github.com/numaproj/numaflow/pkg/daemon/server/service"
 	"github.com/numaproj/numaflow/pkg/isbsvc"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	redisclient "github.com/numaproj/numaflow/pkg/shared/clients/redis"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	sharedtls "github.com/numaproj/numaflow/pkg/shared/tls"
diff --git a/pkg/isb/stores/jetstream/reader.go b/pkg/isb/stores/jetstream/reader.go
index fadf555c..c1d0bbf5 100644
--- a/pkg/isb/stores/jetstream/reader.go
+++ b/pkg/isb/stores/jetstream/reader.go
@@ -30,7 +30,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/wait"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	sharedqueue "github.com/numaproj/numaflow/pkg/shared/queue"
 )
diff --git a/pkg/isb/stores/jetstream/writer.go b/pkg/isb/stores/jetstream/writer.go
index f564bd85..992d593a 100644
--- a/pkg/isb/stores/jetstream/writer.go
+++ b/pkg/isb/stores/jetstream/writer.go
@@ -27,7 +27,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	sharedqueue "github.com/numaproj/numaflow/pkg/shared/queue"
 	sharedutil "github.com/numaproj/numaflow/pkg/shared/util"
diff --git a/pkg/isbsvc/jetstream_service.go b/pkg/isbsvc/jetstream_service.go
index 95fffd1a..8aa2ac0c 100644
--- a/pkg/isbsvc/jetstream_service.go
+++ b/pkg/isbsvc/jetstream_service.go
@@ -24,7 +24,7 @@ import (
 
 	"github.com/nats-io/nats.go"
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/store"
diff --git a/pkg/shared/clients/jetstream/default_jetstream_client.go b/pkg/shared/clients/nats/default_jetstream_client.go
similarity index 99%
rename from pkg/shared/clients/jetstream/default_jetstream_client.go
rename to pkg/shared/clients/nats/default_jetstream_client.go
index 0da05155..bf5628a3 100644
--- a/pkg/shared/clients/jetstream/default_jetstream_client.go
+++ b/pkg/shared/clients/nats/default_jetstream_client.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package jetstream
+package nats
 
 import (
 	"context"
diff --git a/pkg/shared/clients/jetstream/doc.go b/pkg/shared/clients/nats/doc.go
similarity index 98%
rename from pkg/shared/clients/jetstream/doc.go
rename to pkg/shared/clients/nats/doc.go
index 7fdd2eef..10ff88dd 100644
--- a/pkg/shared/clients/jetstream/doc.go
+++ b/pkg/shared/clients/nats/doc.go
@@ -31,4 +31,5 @@ limitations under the License.
 // For example:
 //
 // client.Connect(ctx, AutoReconnect())
-package jetstream
+
+package nats
diff --git a/pkg/shared/clients/jetstream/in_cluster_jetstream_client.go b/pkg/shared/clients/nats/in_cluster_jetstream_client.go
similarity index 99%
rename from pkg/shared/clients/jetstream/in_cluster_jetstream_client.go
rename to pkg/shared/clients/nats/in_cluster_jetstream_client.go
index cd8b031d..9004b617 100644
--- a/pkg/shared/clients/jetstream/in_cluster_jetstream_client.go
+++ b/pkg/shared/clients/nats/in_cluster_jetstream_client.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package jetstream
+package nats
 
 import (
 	"context"
diff --git a/pkg/shared/clients/jetstream/interface.go b/pkg/shared/clients/nats/interface.go
similarity index 97%
rename from pkg/shared/clients/jetstream/interface.go
rename to pkg/shared/clients/nats/interface.go
index b0dda03d..0c174f3e 100644
--- a/pkg/shared/clients/jetstream/interface.go
+++ b/pkg/shared/clients/nats/interface.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package jetstream
+package nats
 
 import "context"
 
diff --git a/pkg/shared/clients/jetstream/jetstream_context.go b/pkg/shared/clients/nats/jetstream_context.go
similarity index 99%
rename from pkg/shared/clients/jetstream/jetstream_context.go
rename to pkg/shared/clients/nats/jetstream_context.go
index 786ac802..08a38a8d 100644
--- a/pkg/shared/clients/jetstream/jetstream_context.go
+++ b/pkg/shared/clients/nats/jetstream_context.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package jetstream
+package nats
 
 import "github.com/nats-io/nats.go"
 
diff --git a/pkg/shared/clients/jetstream/nats_conn.go b/pkg/shared/clients/nats/nats_conn.go
similarity index 99%
rename from pkg/shared/clients/jetstream/nats_conn.go
rename to pkg/shared/clients/nats/nats_conn.go
index bfc81c49..b3ba5ab1 100644
--- a/pkg/shared/clients/jetstream/nats_conn.go
+++ b/pkg/shared/clients/nats/nats_conn.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package jetstream
+package nats
 
 import (
 	"fmt"
diff --git a/pkg/shared/clients/jetstream/options.go b/pkg/shared/clients/nats/options.go
similarity index 99%
rename from pkg/shared/clients/jetstream/options.go
rename to pkg/shared/clients/nats/options.go
index 414941f0..f59f4809 100644
--- a/pkg/shared/clients/jetstream/options.go
+++ b/pkg/shared/clients/nats/options.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package jetstream
+package nats
 
 import (
 	"time"
diff --git a/pkg/sinks/sink.go b/pkg/sinks/sink.go
index 56b39852..f0d19832 100644
--- a/pkg/sinks/sink.go
+++ b/pkg/sinks/sink.go
@@ -31,7 +31,7 @@ import (
 	redisisb "github.com/numaproj/numaflow/pkg/isb/stores/redis"
 	"github.com/numaproj/numaflow/pkg/isbsvc"
 	"github.com/numaproj/numaflow/pkg/metrics"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	redisclient "github.com/numaproj/numaflow/pkg/shared/clients/redis"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	kafkasink "github.com/numaproj/numaflow/pkg/sinks/kafka"
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index 1e08b809..454dfbbd 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -29,7 +29,7 @@ import (
 	redisisb "github.com/numaproj/numaflow/pkg/isb/stores/redis"
 	"github.com/numaproj/numaflow/pkg/isbsvc"
 	"github.com/numaproj/numaflow/pkg/metrics"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	redisclient "github.com/numaproj/numaflow/pkg/shared/clients/redis"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/sources/generator"
diff --git a/pkg/udf/common.go b/pkg/udf/common.go
index 6881e6a3..65e2613a 100644
--- a/pkg/udf/common.go
+++ b/pkg/udf/common.go
@@ -25,7 +25,7 @@ import (
 	jetstreamisb "github.com/numaproj/numaflow/pkg/isb/stores/jetstream"
 	redisisb "github.com/numaproj/numaflow/pkg/isb/stores/redis"
 	"github.com/numaproj/numaflow/pkg/isbsvc"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	redisclient "github.com/numaproj/numaflow/pkg/shared/clients/redis"
 )
 
diff --git a/pkg/watermark/generic/jetstream/generic.go b/pkg/watermark/generic/jetstream/generic.go
index f0b3bce0..9243c3ad 100644
--- a/pkg/watermark/generic/jetstream/generic.go
+++ b/pkg/watermark/generic/jetstream/generic.go
@@ -28,7 +28,7 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isbsvc"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
diff --git a/pkg/watermark/store/jetstream/kv_store.go b/pkg/watermark/store/jetstream/kv_store.go
index 9fd346ea..b0e9c45a 100644
--- a/pkg/watermark/store/jetstream/kv_store.go
+++ b/pkg/watermark/store/jetstream/kv_store.go
@@ -26,7 +26,7 @@ import (
 	"github.com/nats-io/nats.go"
 	"go.uber.org/zap"
 
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
diff --git a/pkg/watermark/store/jetstream/kv_watch.go b/pkg/watermark/store/jetstream/kv_watch.go
index cd4446c1..21658215 100644
--- a/pkg/watermark/store/jetstream/kv_watch.go
+++ b/pkg/watermark/store/jetstream/kv_watch.go
@@ -24,7 +24,7 @@ import (
 	"github.com/nats-io/nats.go"
 	"go.uber.org/zap"
 
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 )


Test patch:
diff --git a/pkg/isb/stores/jetstream/reader_test.go b/pkg/isb/stores/jetstream/reader_test.go
index 568314ef..b3321189 100644
--- a/pkg/isb/stores/jetstream/reader_test.go
+++ b/pkg/isb/stores/jetstream/reader_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
@@ -28,17 +26,18 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
+	natstest "github.com/numaproj/numaflow/pkg/shared/clients/nats/test"
 )
 
-var natsJetStreamUrl = "nats://localhost:4222"
-
 func TestJetStreamBufferRead(t *testing.T) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
+
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancel()
 
-	opts := nats.UserInfo("", "")
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(natsJetStreamUrl, opts)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -64,8 +63,15 @@ func TestJetStreamBufferRead(t *testing.T) {
 		}
 	}
 	// Test Write
-	_, errs := jw.Write(ctx, messages)
+	ofs, errs := jw.Write(ctx, messages)
+	assert.Equal(t, len(ofs), 20)
+	for _, o := range ofs {
+		assert.NotNil(t, o)
+	}
 	assert.Equal(t, len(errs), 20)
+	for _, e := range errs {
+		assert.NoError(t, e)
+	}
 
 	bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName)
 	assert.NoError(t, err)
@@ -74,14 +80,15 @@ func TestJetStreamBufferRead(t *testing.T) {
 
 	readMessages, err := fromStep.Read(ctx, 20)
 	assert.NoError(t, err)
-	offsetsInsideReadMessages := make([]isb.Offset, len(messages))
-	messagesInsideReadMessages := make([]isb.Message, len(messages))
+	assert.Equal(t, 20, len(readMessages))
+	for _, m := range readMessages {
+		assert.NotNil(t, m)
+	}
+	offsetsInsideReadMessages := make([]isb.Offset, len(readMessages))
 	for idx, readMessage := range readMessages {
-		messagesInsideReadMessages[idx] = readMessage.Message
 		offsetsInsideReadMessages[idx] = readMessage.ReadOffset
 	}
-
-	assert.Equal(t, 20, len(readMessages))
+	assert.Equal(t, 20, len(offsetsInsideReadMessages))
 
 	fromStepJs, err := fromStep.conn.JetStream()
 	assert.NoError(t, err)
@@ -116,12 +123,12 @@ func TestJetStreamBufferRead(t *testing.T) {
 
 // TestGetName is used to test the GetName function
 func TestGetName(t *testing.T) {
-
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancel()
 
-	opts := nats.UserInfo("", "")
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(natsJetStreamUrl, opts)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -140,12 +147,13 @@ func TestGetName(t *testing.T) {
 
 // TestClose is used to test Close
 func TestClose(t *testing.T) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancel()
 
-	opts := nats.UserInfo("", "")
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(natsJetStreamUrl, opts)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -174,6 +182,9 @@ func TestConvert2IsbMsgHeader(t *testing.T) {
 }
 
 func addStream(t *testing.T, js *jsclient.JetStreamContext, streamName string) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
+
 	_, err := js.AddStream(&nats.StreamConfig{
 		Name:       streamName,
 		Subjects:   []string{streamName},
diff --git a/pkg/isb/stores/jetstream/writer_test.go b/pkg/isb/stores/jetstream/writer_test.go
index cad8576a..4485825b 100644
--- a/pkg/isb/stores/jetstream/writer_test.go
+++ b/pkg/isb/stores/jetstream/writer_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
@@ -25,14 +23,13 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 
-	"github.com/nats-io/nats.go"
 	"github.com/stretchr/testify/assert"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/forward"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	natstest "github.com/numaproj/numaflow/pkg/shared/clients/nats/test"
 )
 
 type myForwardJetStreamTest struct {
@@ -48,11 +45,12 @@ func (f myForwardJetStreamTest) ApplyMap(ctx context.Context, message *isb.ReadM
 
 // TestForwarderJetStreamBuffer is a test that is used to test forwarder with jetstream buffer
 func TestForwarderJetStreamBuffer(t *testing.T) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
+
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancel()
-
-	opts := nats.UserInfo("", "")
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(natsJetStreamUrl, opts)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -159,11 +157,13 @@ func TestForwarderJetStreamBuffer(t *testing.T) {
 
 // TestJetStreamBufferWrite on buffer full
 func TestJetStreamBufferWriterBufferFull(t *testing.T) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
+
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancel()
 
-	opts := nats.UserInfo("", "")
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(natsJetStreamUrl, opts)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -214,12 +214,13 @@ func TestJetStreamBufferWriterBufferFull(t *testing.T) {
 
 // TestGetName is used to test the GetName function
 func TestWriteGetName(t *testing.T) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancel()
 
-	opts := nats.UserInfo("", "")
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(natsJetStreamUrl, opts)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -239,12 +240,13 @@ func TestWriteGetName(t *testing.T) {
 
 // TestWriteClose is used to test Close function in Write
 func TestWriteClose(t *testing.T) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancel()
 
-	opts := nats.UserInfo("", "")
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(natsJetStreamUrl, opts)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -272,7 +274,5 @@ func TestConvert2NatsMsgHeader(t *testing.T) {
 		},
 		ID: "1",
 	}
-
 	assert.NotNil(t, convert2NatsMsgHeader(isbHeader))
-
 }
diff --git a/pkg/shared/clients/nats/test/client.go b/pkg/shared/clients/nats/test/client.go
new file mode 100644
index 00000000..c6b68933
--- /dev/null
+++ b/pkg/shared/clients/nats/test/client.go
@@ -0,0 +1,30 @@
+/*
+Copyright 2022 The Numaproj 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 (
+	"testing"
+
+	"github.com/nats-io/nats-server/v2/server"
+
+	"github.com/numaproj/numaflow/pkg/shared/clients/nats"
+)
+
+// JetStreamClient is used to get a testing JetStream client instance
+func JetStreamClient(t *testing.T, s *server.Server) nats.JetStreamClient {
+	return nats.NewDefaultJetStreamClient(s.ClientURL())
+}
diff --git a/pkg/shared/clients/nats/test/server.go b/pkg/shared/clients/nats/test/server.go
new file mode 100644
index 00000000..1f2a4d5f
--- /dev/null
+++ b/pkg/shared/clients/nats/test/server.go
@@ -0,0 +1,62 @@
+/*
+Copyright 2022 The Numaproj 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 (
+	"os"
+	"testing"
+
+	"github.com/nats-io/nats-server/v2/server"
+	natstestserver "github.com/nats-io/nats-server/v2/test"
+)
+
+// RunNatsServer starts a nats server
+func RunNatsServer(t *testing.T) *server.Server {
+	t.Helper()
+	opts := natstestserver.DefaultTestOptions
+	return natstestserver.RunServer(&opts)
+}
+
+// RunJetStreamServer starts a jetstream server
+func RunJetStreamServer(t *testing.T) *server.Server {
+	t.Helper()
+	opts := natstestserver.DefaultTestOptions
+	opts.Port = -1 // Random port
+	opts.JetStream = true
+	storeDir, err := os.MkdirTemp("", "")
+	if err != nil {
+		t.Fatalf("Error creating a temp dir: %v", err)
+	}
+	opts.StoreDir = storeDir
+	return natstestserver.RunServer(&opts)
+}
+
+// ShutdownJetStreamServer shuts down the jetstream server and clean up resources
+func ShutdownJetStreamServer(t *testing.T, s *server.Server) {
+	t.Helper()
+	var sd string
+	if config := s.JetStreamConfig(); config != nil {
+		sd = config.StoreDir
+	}
+	s.Shutdown()
+	if sd != "" {
+		if err := os.RemoveAll(sd); err != nil {
+			t.Fatalf("Failed to remove storage %q: %v", sd, err)
+		}
+	}
+	s.WaitForShutdown()
+}
diff --git a/pkg/sources/nats/nats_test.go b/pkg/sources/nats/nats_test.go
index d456631c..52c80418 100644
--- a/pkg/sources/nats/nats_test.go
+++ b/pkg/sources/nats/nats_test.go
@@ -22,14 +22,13 @@ import (
 	"testing"
 	"time"
 
-	"github.com/nats-io/nats-server/v2/server"
-	natstestserver "github.com/nats-io/nats-server/v2/test"
 	natslib "github.com/nats-io/nats.go"
 	"github.com/stretchr/testify/assert"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
+	natstest "github.com/numaproj/numaflow/pkg/shared/clients/nats/test"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/noop"
@@ -67,19 +66,11 @@ func newInstance(t *testing.T, vi *dfv1.VertexInstance) (*natsSource, error) {
 	return New(vi, dest, fetchWatermark, publishWatermark, publishWMStores, WithReadTimeout(1*time.Second))
 }
 
-func runNatsServer(t *testing.T) *server.Server {
-	t.Helper()
-	opts := natstestserver.DefaultTestOptions
-	opts.Port = 9989
-	opts.Cluster.Name = "testing"
-	return natstestserver.RunServer(&opts)
-}
-
 func Test_Single(t *testing.T) {
-	server := runNatsServer(t)
+	server := natstest.RunNatsServer(t)
 	defer server.Shutdown()
 
-	url := "127.0.0.1:9989"
+	url := "127.0.0.1"
 	testSubject := "test"
 	testQueue := "test-queue"
 	vi := testVertex(t, url, testSubject, testQueue, "test-host", 0)
@@ -102,10 +93,10 @@ func Test_Single(t *testing.T) {
 }
 
 func Test_Multiple(t *testing.T) {
-	server := runNatsServer(t)
+	server := natstest.RunNatsServer(t)
 	defer server.Shutdown()
 
-	url := "127.0.0.1:9989"
+	url := "127.0.0.1"
 	testSubject := "test"
 	testQueue := "test-queue"
 	v1 := testVertex(t, url, testSubject, testQueue, "test-host1", 0)
diff --git a/pkg/watermark/fetch/offset_timeline_test.go b/pkg/watermark/fetch/offset_timeline_test.go
index eafb2a9e..7f370d24 100644
--- a/pkg/watermark/fetch/offset_timeline_test.go
+++ b/pkg/watermark/fetch/offset_timeline_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
diff --git a/pkg/watermark/fetch/processor_heartbeat_test.go b/pkg/watermark/fetch/processor_heartbeat_test.go
index 778d3a11..3194b138 100644
--- a/pkg/watermark/fetch/processor_heartbeat_test.go
+++ b/pkg/watermark/fetch/processor_heartbeat_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
diff --git a/pkg/watermark/fetch/processor_manager_test.go b/pkg/watermark/fetch/processor_manager_test.go
index 01ddb509..ee13d9fa 100644
--- a/pkg/watermark/fetch/processor_manager_test.go
+++ b/pkg/watermark/fetch/processor_manager_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
@@ -30,7 +28,7 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	natstest "github.com/numaproj/numaflow/pkg/shared/clients/nats/test"
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
 )
@@ -42,11 +40,15 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 		testOffset int64 = 100
 		wg         sync.WaitGroup
 	)
+
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
+
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
 	defer cancel()
 
 	// connect to NATS
-	nc, err := jsclient.NewDefaultJetStreamClient(nats.DefaultURL).Connect(context.TODO())
+	nc, err := natstest.JetStreamClient(t, s).Connect(context.TODO())
 	assert.NoError(t, err)
 	defer nc.Close()
 
@@ -84,8 +86,7 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 	defer func() { _ = js.DeleteKeyValue(keyspace + "_OT") }()
 	assert.NoError(t, err)
 
-	// create default JetStream client
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(nats.DefaultURL)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 
 	// create hbStore
 	hbStore, err := jetstream.NewKVJetStreamKVStore(ctx, "testFetch", keyspace+"_PROCESSORS", defaultJetStreamClient)
diff --git a/pkg/watermark/fetch/processor_to_fetch_test.go b/pkg/watermark/fetch/processor_to_fetch_test.go
index 6ec32e02..675a6343 100644
--- a/pkg/watermark/fetch/processor_to_fetch_test.go
+++ b/pkg/watermark/fetch/processor_to_fetch_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
diff --git a/pkg/watermark/fetch/source_fetcher_test.go b/pkg/watermark/fetch/source_fetcher_test.go
index b7c0da75..0c39635f 100644
--- a/pkg/watermark/fetch/source_fetcher_test.go
+++ b/pkg/watermark/fetch/source_fetcher_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
diff --git a/pkg/watermark/publish/publisher_test.go b/pkg/watermark/publish/publisher_test.go
index 128807c9..013052c7 100644
--- a/pkg/watermark/publish/publisher_test.go
+++ b/pkg/watermark/publish/publisher_test.go
@@ -1,5 +1,3 @@
-//go:build isb_jetstream
-
 /*
 Copyright 2022 The Numaproj Authors.
 
@@ -28,7 +26,8 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/jetstream"
+	jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats"
+	natstest "github.com/numaproj/numaflow/pkg/shared/clients/nats/test"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/numaproj/numaflow/pkg/watermark/store/jetstream"
@@ -42,14 +41,16 @@ func createAndLaterDeleteBucket(js *jsclient.JetStreamContext, kvConfig *nats.Ke
 	return func() {
 		// ignore error as this will be executed only via defer
 		_ = js.DeleteKeyValue(kv.Bucket())
-		return
 	}, nil
 }
 
 func TestPublisherWithSharedOTBucket(t *testing.T) {
+	s := natstest.RunJetStreamServer(t)
+	defer natstest.ShutdownJetStreamServer(t, s)
+
 	var ctx = context.Background()
 
-	defaultJetStreamClient := jsclient.NewDefaultJetStreamClient(nats.DefaultURL)
+	defaultJetStreamClient := natstest.JetStreamClient(t, s)
 	conn, err := defaultJetStreamClient.Connect(ctx)
 	assert.NoError(t, err)
 	js, err := conn.JetStream()
@@ -58,9 +59,11 @@ func TestPublisherWithSharedOTBucket(t *testing.T) {
 	var keyspace = "publisherTest"
 
 	deleteFn, err := createAndLaterDeleteBucket(js, &nats.KeyValueConfig{Bucket: keyspace + "_PROCESSORS"})
+	assert.NoError(t, err)
 	defer deleteFn()
 
 	deleteFn, err = createAndLaterDeleteBucket(js, &nats.KeyValueConfig{Bucket: keyspace + "_OT"})
+	assert.NoError(t, err)
 	defer deleteFn()
 
 	publishEntity := processor.NewProcessorEntity("publisherTestPod1")