Repository: gofr-dev/gofr

Issue:
Report a bug of connect to NATS
**Describe the bug**
The current code uses an asynchronous way to connect to the NATS server：
```
func (cm *ConnectionManager) Connect() error {
	go cm.retryConnect()

	return nil
}

```
In Client.Connect() ，calling connManager.Connect() and then immediately getting jetStream() is likely to fail, resulting in the errJetStreamNotConfigured , because the asynchronous connection has not been completed yet.
```
      if err := connManager.Connect(); err != nil {
		c.logger.Errorf("failed to connect to NATS server at %v: %v", c.Config.Server, err)
		return err
	}

	c.connManager = connManager

	js, err := c.connManager.jetStream()
	if err != nil {
		return err
	}
```
**To Reproduce**
Steps to reproduce the behavior, if applicable:

1. The code is

   ```go
   app := gofr.New()
    app.AddPubSub(nats.New(nats.Config{
        Server:     "nats://localhost:4222",
        Stream: nats.StreamConfig{
            Stream:   "mystream",
            Subjects: []string{"orders.*", "shipments.*"},
        },
        MaxWait:     5 * time.Second,
        MaxPullWait: 500 * time.Millisecond,
        Consumer:    "my-consumer",
        CredsFile:   "/path/to/creds.json",
    }))
      ...
    app.Run()
  ```

2. The error is

   ```
   
   ```

**Expected behavior**
A clear and concise description of what you expected to happen.

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environments (please complete the following information):**
- OS: [e.g. Linux]
- gofr version [e.g. v1.5.0]
- go version [e.g. 1.21]

**More description**
Add any other context about the problem here.


Hints from issue comments:
@Sumit-Kumar1: Hi @Umang01-hash, please assign this to me.@Umang01-hash: Sure @Sumit-Kumar1 done!

Solution patch:
diff --git a/pkg/gofr/datasource/pubsub/nats/client.go b/pkg/gofr/datasource/pubsub/nats/client.go
index 88e2fe52..4028e0ea 100644
--- a/pkg/gofr/datasource/pubsub/nats/client.go
+++ b/pkg/gofr/datasource/pubsub/nats/client.go
@@ -15,7 +15,7 @@ import (
 
 //go:generate mockgen -destination=mock_tracer.go -package=nats go.opentelemetry.io/otel/trace Tracer
 
-const defaultRetryTimeout = 10 * time.Second
+const defaultRetryTimeout = 2 * time.Second
 
 var errClientNotConnected = errors.New("nats client not connected")
 
@@ -44,9 +44,21 @@ func (c *Client) Connect() error {
 		return err
 	}
 
+	if err := c.establishConnection(); err != nil {
+		c.logger.Errorf("failed to connect to NATS server at %v: %v", c.Config.Server, err)
+
+		go c.retryConnect()
+
+		return err
+	}
+
+	return nil
+}
+
+// establishConnection handles the actual connection process to NATS and sets up jStream.
+func (c *Client) establishConnection() error {
 	connManager := NewConnectionManager(c.Config, c.logger, c.natsConnector, c.jetStreamCreator)
 	if err := connManager.Connect(); err != nil {
-		c.logger.Errorf("failed to connect to NATS server at %v: %v", c.Config.Server, err)
 		return err
 	}
 
@@ -64,6 +76,22 @@ func (c *Client) Connect() error {
 	return nil
 }
 
+func (c *Client) retryConnect() {
+	for {
+		c.logger.Debugf("connecting to NATS server at %v", c.Config.Server)
+
+		if err := c.establishConnection(); err != nil {
+			c.logger.Errorf("failed to connect to NATS server at %v: %v", c.Config.Server, err)
+
+			time.Sleep(defaultRetryTimeout)
+
+			continue
+		}
+
+		return
+	}
+}
+
 func (c *Client) validateAndPrepare() error {
 	if err := validateConfigs(c.Config); err != nil {
 		c.logger.Errorf("could not initialize NATS jStream: %v", err)
@@ -103,13 +131,17 @@ func (c *Client) UseMetrics(metrics any) {
 
 // Publish publishes a message to a topic.
 func (c *Client) Publish(ctx context.Context, subject string, message []byte) error {
+	if err := checkClient(c); err != nil {
+		return err
+	}
+
 	return c.connManager.Publish(ctx, subject, message, c.metrics)
 }
 
 // Subscribe subscribes to a topic and returns a single message.
 func (c *Client) Subscribe(ctx context.Context, topic string) (*pubsub.Message, error) {
 	for {
-		if !c.connManager.isConnected() {
+		if err := checkClient(c); err != nil {
 			time.Sleep(defaultRetryTimeout)
 
 			return nil, errClientNotConnected
@@ -256,6 +288,10 @@ func (c *Client) Close(ctx context.Context) error {
 
 // CreateTopic creates a new topic (stream) in NATS jStream.
 func (c *Client) CreateTopic(ctx context.Context, name string) error {
+	if err := checkClient(c); err != nil {
+		return err
+	}
+
 	return c.streamManager.CreateStream(ctx, StreamConfig{
 		Stream:   name,
 		Subjects: []string{name},
@@ -264,21 +300,37 @@ func (c *Client) CreateTopic(ctx context.Context, name string) error {
 
 // DeleteTopic deletes a topic (stream) in NATS jStream.
 func (c *Client) DeleteTopic(ctx context.Context, name string) error {
+	if err := checkClient(c); err != nil {
+		return err
+	}
+
 	return c.streamManager.DeleteStream(ctx, name)
 }
 
 // CreateStream creates a new stream in NATS jStream.
 func (c *Client) CreateStream(ctx context.Context, cfg StreamConfig) error {
+	if err := checkClient(c); err != nil {
+		return err
+	}
+
 	return c.streamManager.CreateStream(ctx, cfg)
 }
 
 // DeleteStream deletes a stream in NATS jStream.
 func (c *Client) DeleteStream(ctx context.Context, name string) error {
+	if err := checkClient(c); err != nil {
+		return err
+	}
+
 	return c.streamManager.DeleteStream(ctx, name)
 }
 
 // CreateOrUpdateStream creates or updates a stream in NATS jStream.
 func (c *Client) CreateOrUpdateStream(ctx context.Context, cfg *jetstream.StreamConfig) (jetstream.Stream, error) {
+	if err := checkClient(c); err != nil {
+		return nil, err
+	}
+
 	return c.streamManager.CreateOrUpdateStream(ctx, cfg)
 }
 
@@ -291,3 +343,19 @@ func GetJetStreamStatus(ctx context.Context, js jetstream.JetStream) (string, er
 
 	return jetStreamStatusOK, nil
 }
+
+func checkClient(c *Client) error {
+	if c == nil {
+		return errClientNotConnected
+	}
+
+	if c.connManager == nil {
+		return errClientNotConnected
+	}
+
+	if !c.connManager.isConnected() {
+		return errClientNotConnected
+	}
+
+	return nil
+}
diff --git a/pkg/gofr/datasource/pubsub/nats/connection_manager.go b/pkg/gofr/datasource/pubsub/nats/connection_manager.go
index 3cc32592..a3d93d3e 100644
--- a/pkg/gofr/datasource/pubsub/nats/connection_manager.go
+++ b/pkg/gofr/datasource/pubsub/nats/connection_manager.go
@@ -2,7 +2,6 @@ package nats
 
 import (
 	"context"
-	"time"
 
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
@@ -79,42 +78,31 @@ func NewConnectionManager(
 
 // Connect establishes a connection to NATS and sets up JetStream.
 func (cm *ConnectionManager) Connect() error {
-	go cm.retryConnect()
-
-	return nil
-}
-
-func (cm *ConnectionManager) retryConnect() {
 	opts := []nats.Option{nats.Name("GoFr NATS JetStreamClient")}
 
 	if cm.config.CredsFile != "" {
 		opts = append(opts, nats.UserCredentials(cm.config.CredsFile))
 	}
 
-	for {
-		connInterface, err := cm.natsConnector.Connect(cm.config.Server, opts...)
-		if err != nil {
-			cm.logger.Errorf("Failed to connect to NATS server at %v: %v", cm.config.Server, err)
-			time.Sleep(defaultRetryTimeout)
-
-			continue
-		}
+	connInterface, err := cm.natsConnector.Connect(cm.config.Server, opts...)
+	if err != nil {
+		cm.logger.Debugf("Failed to connect to NATS server at %v: %v", cm.config.Server, err)
+		return err
+	}
 
-		js, err := cm.jetStreamCreator.New(connInterface)
-		if err != nil {
-			connInterface.Close()
-			cm.logger.Debugf("Failed to create jStream context: %v", err)
-			time.Sleep(defaultRetryTimeout)
+	js, err := cm.jetStreamCreator.New(connInterface)
+	if err != nil {
+		connInterface.Close()
+		cm.logger.Debugf("Failed to create jStream context: %v", err)
 
-			continue
-		}
+		return err
+	}
 
-		cm.conn = connInterface
-		cm.jStream = js
-		cm.logger.Logf("Successfully connected to NATS server at %v", cm.config.Server)
+	cm.conn = connInterface
+	cm.jStream = js
+	cm.logger.Logf("Successfully connected to NATS server at %v", cm.config.Server)
 
-		return
-	}
+	return nil
 }
 
 func (cm *ConnectionManager) Close(_ context.Context) {
@@ -126,6 +114,10 @@ func (cm *ConnectionManager) Close(_ context.Context) {
 func (cm *ConnectionManager) Publish(ctx context.Context, subject string, message []byte, metrics Metrics) error {
 	metrics.IncrementCounter(ctx, "app_pubsub_publish_total_count", "subject", subject)
 
+	if !cm.isConnected() {
+		return errClientNotConnected
+	}
+
 	if err := cm.validateJetStream(subject); err != nil {
 		return err
 	}



Test patch:
diff --git a/pkg/gofr/datasource/pubsub/nats/client_test.go b/pkg/gofr/datasource/pubsub/nats/client_test.go
index 18d3c1e7..7a1228bc 100644
--- a/pkg/gofr/datasource/pubsub/nats/client_test.go
+++ b/pkg/gofr/datasource/pubsub/nats/client_test.go
@@ -49,9 +49,10 @@ func TestNATSClient_Publish(t *testing.T) {
 	message := []byte("test-message")
 
 	// Set up expected calls
-	mockConnManager.EXPECT().
-		Publish(ctx, subject, message, mockMetrics).
-		Return(nil)
+	gomock.InOrder(
+		mockConnManager.EXPECT().IsConnected().Return(true),
+		mockConnManager.EXPECT().Publish(ctx, subject, message, mockMetrics).Return(nil),
+	)
 
 	// Call Publish
 	err := client.Publish(ctx, subject, message)
@@ -65,6 +66,10 @@ func TestNATSClient_PublishError(t *testing.T) {
 	mockMetrics := NewMockMetrics(ctrl)
 	mockConnManager := NewMockConnectionManagerInterface(ctrl)
 
+	ctx := context.TODO()
+	subject := "test"
+	message := []byte("test-message")
+
 	config := &Config{
 		Server: NATSServer,
 		Stream: StreamConfig{
@@ -81,18 +86,35 @@ func TestNATSClient_PublishError(t *testing.T) {
 		logger:      logging.NewMockLogger(logging.DEBUG),
 	}
 
-	ctx := context.TODO()
-	subject := "test"
-	message := []byte("test-message")
+	tests := []struct {
+		name     string
+		client   *Client
+		mockCall func(mockConn *MockConnectionManagerInterface)
+		expErr   error
+	}{
+		{name: "nil client", client: nil, mockCall: func(_ *MockConnectionManagerInterface) {}, expErr: errClientNotConnected},
+		{name: "nil connManager", client: &Client{connManager: nil}, mockCall: func(_ *MockConnectionManagerInterface) {},
+			expErr: errClientNotConnected},
+		{name: "not connected to NATS server", client: &Client{connManager: mockConnManager},
+			mockCall: func(mockConn *MockConnectionManagerInterface) {
+				mockConn.EXPECT().IsConnected().Return(false)
+			}, expErr: errClientNotConnected},
+		{name: "err in publishing", client: client, mockCall: func(mockConn *MockConnectionManagerInterface) {
+			mockConn.EXPECT().IsConnected().Return(true)
+			mockConn.EXPECT().Publish(gomock.Any(), subject, message, mockMetrics).Return(errPublishError)
+		}, expErr: errPublishError},
+	}
 
-	expectedErr := errPublishError
-	mockConnManager.EXPECT().
-		Publish(ctx, subject, message, mockMetrics).
-		Return(expectedErr)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.mockCall(mockConnManager)
 
-	err := client.Publish(ctx, subject, message)
-	require.Error(t, err)
-	assert.Equal(t, expectedErr, err)
+			err := tt.client.Publish(ctx, subject, message)
+
+			require.Error(t, err)
+			assert.Equal(t, tt.expErr, err)
+		})
+	}
 }
 
 func TestNATSClient_SubscribeSuccess(t *testing.T) {
@@ -251,15 +273,20 @@ func TestNATSClient_DeleteTopic(t *testing.T) {
 	defer ctrl.Finish()
 
 	mockStreamManager := NewMockStreamManagerInterface(ctrl)
+	mockConnManager := NewMockConnectionManagerInterface(ctrl)
+
 	client := &Client{
 		streamManager: mockStreamManager,
+		connManager:   mockConnManager,
 		logger:        logging.NewMockLogger(logging.DEBUG),
 		Config:        &Config{},
 	}
 
 	ctx := context.Background()
-
-	mockStreamManager.EXPECT().DeleteStream(ctx, "test-topic").Return(nil)
+	gomock.InOrder(
+		mockConnManager.EXPECT().IsConnected().Return(true),
+		mockStreamManager.EXPECT().DeleteStream(ctx, "test-topic").Return(nil),
+	)
 
 	err := client.DeleteTopic(ctx, "test-topic")
 	require.NoError(t, err)
@@ -270,8 +297,10 @@ func TestNATSClient_DeleteTopic_Error(t *testing.T) {
 	defer ctrl.Finish()
 
 	mockStreamManager := NewMockStreamManagerInterface(ctrl)
+	mockConnManager := NewMockConnectionManagerInterface(ctrl)
 	client := &Client{
 		streamManager: mockStreamManager,
+		connManager:   mockConnManager,
 		logger:        logging.NewMockLogger(logging.DEBUG),
 		Config:        &Config{},
 	}
@@ -279,7 +308,11 @@ func TestNATSClient_DeleteTopic_Error(t *testing.T) {
 	ctx := context.Background()
 
 	expectedErr := errFailedToDeleteStream
-	mockStreamManager.EXPECT().DeleteStream(ctx, "test-topic").Return(expectedErr)
+
+	gomock.InOrder(
+		mockConnManager.EXPECT().IsConnected().Return(true),
+		mockStreamManager.EXPECT().DeleteStream(ctx, "test-topic").Return(expectedErr),
+	)
 
 	err := client.DeleteTopic(ctx, "test-topic")
 	require.Error(t, err)
@@ -291,20 +324,22 @@ func TestNATSClient_CreateTopic(t *testing.T) {
 	defer ctrl.Finish()
 
 	mockStreamManager := NewMockStreamManagerInterface(ctrl)
+	mockConnManager := NewMockConnectionManagerInterface(ctrl)
 	client := &Client{
 		streamManager: mockStreamManager,
+		connManager:   mockConnManager,
 		logger:        logging.NewMockLogger(logging.DEBUG),
 		Config:        &Config{},
 	}
 
 	ctx := context.Background()
 
+	mockConnManager.EXPECT().IsConnected().Return(true)
 	mockStreamManager.EXPECT().
 		CreateStream(ctx, StreamConfig{
 			Stream:   "test-topic",
 			Subjects: []string{"test-topic"},
-		}).
-		Return(nil)
+		}).Return(nil)
 
 	err := client.CreateTopic(ctx, "test-topic")
 	require.NoError(t, err)
@@ -362,14 +397,16 @@ func TestClient_Connect(t *testing.T) {
 		"Successfully connected to NATS server at nats://localhost:4222\n")
 }
 
-func TestClient_ConnectError(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
+func TestClient_ValidateAndPrepare(t *testing.T) {
+	client := &Client{
+		Config: &Config{},
+		logger: logging.NewMockLogger(logging.DEBUG),
+	}
 
-	mockNATSConnector := NewMockNATSConnector(ctrl)
-	mockJSCreator := NewMockJetStreamCreator(ctrl)
+	err := client.validateAndPrepare()
+	require.Error(t, err)
 
-	config := &Config{
+	client.Config = &Config{
 		Server: "nats://localhost:4222",
 		Stream: StreamConfig{
 			Stream:   "test-stream",
@@ -378,38 +415,11 @@ func TestClient_ConnectError(t *testing.T) {
 		Consumer: "test-consumer",
 	}
 
-	client := &Client{
-		Config:           config,
-		logger:           logging.NewMockLogger(logging.DEBUG),
-		natsConnector:    mockNATSConnector,
-		jetStreamCreator: mockJSCreator,
-	}
-
-	// Simulate a connection error
-	expectedErr := errJetStreamNotConfigured
-	mockNATSConnector.EXPECT().
-		Connect(config.Server, gomock.Any()).
-		Return(nil, expectedErr)
-
-	// Capture stderr output
-	output := testutil.StderrOutputForFunc(func() {
-		client.logger = logging.NewMockLogger(logging.DEBUG)
-		err := client.Connect()
-
-		time.Sleep(100 * time.Millisecond)
-
-		require.Error(t, err)
-		assert.Equal(t, expectedErr, err)
-	})
-
-	// Check for the error log
-	assert.Contains(t, output, "Failed to connect to NATS server at nats://localhost:4222: jStream is not configured\n")
-
-	assert.Nil(t, client.streamManager)
-	assert.Nil(t, client.subManager)
+	err = client.validateAndPrepare()
+	assert.NoError(t, err)
 }
 
-func TestClient_ValidateAndPrepare(t *testing.T) {
+func TestClient_ValidateAndPrepareError(t *testing.T) {
 	client := &Client{
 		Config: &Config{},
 		logger: logging.NewMockLogger(logging.DEBUG),
@@ -419,7 +429,7 @@ func TestClient_ValidateAndPrepare(t *testing.T) {
 	require.Error(t, err)
 
 	client.Config = &Config{
-		Server: "nats://localhost:4222",
+		Server: "",
 		Stream: StreamConfig{
 			Stream:   "test-stream",
 			Subjects: []string{"test-subject"},
@@ -428,7 +438,7 @@ func TestClient_ValidateAndPrepare(t *testing.T) {
 	}
 
 	err = client.validateAndPrepare()
-	assert.NoError(t, err)
+	assert.Error(t, err)
 }
 
 func TestClient_LogSuccessfulConnection(t *testing.T) {
@@ -728,7 +738,9 @@ func TestClient_CreateStream(t *testing.T) {
 	defer ctrl.Finish()
 
 	mockStreamManager := NewMockStreamManagerInterface(ctrl)
+	mockConnManager := NewMockConnectionManagerInterface(ctrl)
 	client := &Client{
+		connManager:   mockConnManager,
 		streamManager: mockStreamManager,
 	}
 
@@ -737,6 +749,7 @@ func TestClient_CreateStream(t *testing.T) {
 		Subjects: []string{"test-subject"},
 	}
 
+	mockConnManager.EXPECT().IsConnected().Return(true)
 	mockStreamManager.EXPECT().CreateStream(gomock.Any(), cfg).Return(nil)
 
 	err := client.CreateStream(context.Background(), cfg)
@@ -748,10 +761,13 @@ func TestClient_DeleteStream(t *testing.T) {
 	defer ctrl.Finish()
 
 	mockStreamManager := NewMockStreamManagerInterface(ctrl)
+	mockConnManager := NewMockConnectionManagerInterface(ctrl)
 	client := &Client{
+		connManager:   mockConnManager,
 		streamManager: mockStreamManager,
 	}
 
+	mockConnManager.EXPECT().IsConnected().Return(true)
 	mockStreamManager.EXPECT().DeleteStream(gomock.Any(), "test-stream").Return(nil)
 
 	err := client.DeleteStream(context.Background(), "test-stream")
@@ -763,8 +779,10 @@ func TestClient_CreateOrUpdateStream(t *testing.T) {
 	defer ctrl.Finish()
 
 	mockStreamManager := NewMockStreamManagerInterface(ctrl)
+	mockConnManager := NewMockConnectionManagerInterface(ctrl)
 	mockStream := NewMockStream(ctrl)
 	client := &Client{
+		connManager:   mockConnManager,
 		streamManager: mockStreamManager,
 	}
 
@@ -773,9 +791,220 @@ func TestClient_CreateOrUpdateStream(t *testing.T) {
 		Subjects: []string{"test-subject"},
 	}
 
+	mockConnManager.EXPECT().IsConnected().Return(true)
 	mockStreamManager.EXPECT().CreateOrUpdateStream(gomock.Any(), &cfg).Return(mockStream, nil)
 
 	stream, err := client.CreateOrUpdateStream(context.Background(), &cfg)
 	require.NoError(t, err)
 	assert.Equal(t, mockStream, stream)
 }
+
+func Test_checkClient(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockConn := NewMockConnectionManagerInterface(ctrl)
+
+	tests := []struct {
+		name     string
+		client   *Client
+		mockCall *gomock.Call
+		expErr   error
+	}{
+		{name: "nil client", client: nil, expErr: errClientNotConnected},
+		{name: "nil connManager", client: &Client{connManager: nil}, expErr: errClientNotConnected},
+		{name: "not connected to NATS server", client: &Client{connManager: mockConn},
+			mockCall: mockConn.EXPECT().IsConnected().Return(false), expErr: errClientNotConnected},
+		{name: "valid client", client: &Client{connManager: mockConn},
+			mockCall: mockConn.EXPECT().IsConnected().Return(true), expErr: nil},
+	}
+
+	for i, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := checkClient(tt.client)
+
+			assert.Equalf(t, tt.expErr, err, "Test[%d] failed - %s", i, tt.name)
+		})
+	}
+}
+
+func TestClient_retryConnect(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockNATSConnector := NewMockNATSConnector(ctrl)
+	mockJSCreator := NewMockJetStreamCreator(ctrl)
+	mockConn := NewMockConnInterface(ctrl)
+	mockJS := NewMockJetStream(ctrl)
+	logger := logging.NewMockLogger(logging.DEBUG)
+	subs := make(map[string]context.CancelFunc)
+	cfg := Config{Server: "nats://localhost:4222",
+		Stream:   StreamConfig{Stream: "test_stream", Subjects: []string{"test_subject"}},
+		Consumer: "test_consumer",
+	}
+
+	tests := []struct {
+		name        string
+		setupMocks  func(*Client, *MockNATSConnector, *MockJetStreamCreator, *MockConnInterface, *MockJetStream)
+		connSuccess bool
+	}{
+		{
+			name: "successful connection on first attempt",
+			setupMocks: func(client *Client, mockNATSConnector *MockNATSConnector, mockJSCreator *MockJetStreamCreator,
+				mockConn *MockConnInterface, mockJS *MockJetStream) {
+				gomock.InOrder(
+					mockNATSConnector.EXPECT().Connect(client.Config.Server, gomock.Any()).
+						Return(mockConn, nil).MaxTimes(1),
+					mockJSCreator.EXPECT().New(mockConn).Return(mockJS, nil).MaxTimes(1),
+				)
+			},
+			connSuccess: true,
+		},
+		{
+			name: "successful connection after retries",
+			setupMocks: func(client *Client, mockNATSConnector *MockNATSConnector, mockJSCreator *MockJetStreamCreator,
+				mockConn *MockConnInterface, mockJS *MockJetStream) {
+				gomock.InOrder(
+					// First attempt fails
+					mockNATSConnector.EXPECT().Connect(client.Config.Server, gomock.Any()).
+						Return(nil, errConnectionError).MaxTimes(1),
+					// Second attempt succeeds
+					mockNATSConnector.EXPECT().Connect(client.Config.Server, gomock.Any()).
+						Return(mockConn, nil).MaxTimes(1),
+					mockJSCreator.EXPECT().New(mockConn).
+						Return(mockJS, nil),
+				)
+			},
+		},
+		{
+			name: "JetStream creation fails after successful connection",
+			setupMocks: func(client *Client, mockNATSConnector *MockNATSConnector, mockJSCreator *MockJetStreamCreator,
+				mockConn *MockConnInterface, mockJS *MockJetStream) {
+				gomock.InOrder(
+					// Connection succeeds but JetStream creation fails
+					mockNATSConnector.EXPECT().Connect(client.Config.Server, gomock.Any()).Return(mockConn, nil),
+					mockJSCreator.EXPECT().New(mockConn).Return(nil, errConnectionError),
+					mockConn.EXPECT().Close(),
+					// Retry succeeds
+					mockNATSConnector.EXPECT().Connect(client.Config.Server, gomock.Any()).Return(mockConn, nil),
+					mockJSCreator.EXPECT().New(mockConn).Return(mockJS, nil),
+				)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		client := &Client{
+			Config:           &cfg,
+			logger:           logger,
+			natsConnector:    mockNATSConnector,
+			jetStreamCreator: mockJSCreator,
+			subscriptions:    subs,
+		}
+
+		tt.setupMocks(client, mockNATSConnector, mockJSCreator, mockConn, mockJS)
+		client.retryConnect()
+		time.Sleep(500 * time.Millisecond)
+
+		if tt.connSuccess {
+			assert.NotNil(t, client.connManager)
+			assert.NotNil(t, client.streamManager)
+			assert.NotNil(t, client.subManager)
+		}
+	}
+}
+
+func TestClient_GetJetStreamStatus(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	jStream := NewMockJetStream(ctrl)
+
+	tests := []struct {
+		name     string
+		mockCall *gomock.Call
+		want     string
+		wantErr  error
+	}{
+		{name: "status OK", want: jetStreamStatusOK,
+			mockCall: jStream.EXPECT().AccountInfo(gomock.Any()).Return(nil, nil)},
+		{name: "error in jetstream", want: jetStreamStatusError, wantErr: errJetStream,
+			mockCall: jStream.EXPECT().AccountInfo(gomock.Any()).Return(nil, errJetStream)},
+	}
+
+	for i, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+
+			got, err := GetJetStreamStatus(ctx, jStream)
+
+			assert.Equalf(t, tt.wantErr, err, "Test[%d] failed- %s", i, tt.name)
+			assert.Equalf(t, tt.want, got, "Test[%d] failed- %s", i, tt.name)
+		})
+	}
+}
+
+func TestClient_establishConnection(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockNATSConnector := NewMockNATSConnector(ctrl)
+	mockJSCreator := NewMockJetStreamCreator(ctrl)
+	mockConn := NewMockConnInterface(ctrl)
+	mockJS := NewMockJetStream(ctrl)
+	logger := logging.NewMockLogger(logging.DEBUG)
+	cfg := Config{
+		Server: "nats://localhost:4222",
+		Stream: StreamConfig{
+			Stream:   "test_stream",
+			Subjects: []string{"test_subject"},
+		},
+		Consumer: "test_consumer",
+	}
+
+	tests := []struct {
+		name       string
+		client     *Client
+		setupMocks func(*Client, *MockNATSConnector, *MockJetStreamCreator, *MockConnInterface, *MockJetStream)
+		wantErr    error
+	}{
+		{
+			name: "successful connection",
+			setupMocks: func(client *Client, mockNATSConnector *MockNATSConnector, _ *MockJetStreamCreator,
+				_ *MockConnInterface, mockJS *MockJetStream) {
+				gomock.InOrder(
+					mockNATSConnector.EXPECT().Connect(client.Config.Server, gomock.Any()).
+						Return(mockConn, nil),
+					mockJSCreator.EXPECT().New(mockConn).Return(mockJS, nil),
+				)
+			},
+			wantErr: nil,
+		},
+		{
+			name: "connection failure",
+			setupMocks: func(client *Client, mockNATSConnector *MockNATSConnector, _ *MockJetStreamCreator,
+				_ *MockConnInterface, _ *MockJetStream) {
+				mockNATSConnector.EXPECT().Connect(client.Config.Server, gomock.Any()).
+					Return(nil, errConnectionError)
+			},
+			wantErr: errConnectionError,
+		},
+	}
+
+	for i, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			client := &Client{
+				Config:           &cfg,
+				logger:           logger,
+				natsConnector:    mockNATSConnector,
+				jetStreamCreator: mockJSCreator,
+			}
+
+			tt.setupMocks(client, mockNATSConnector, mockJSCreator, mockConn, mockJS)
+
+			err := client.establishConnection()
+
+			assert.Equalf(t, tt.wantErr, err, "Test[%d] failed - %s", i, tt.name)
+		})
+	}
+}
diff --git a/pkg/gofr/datasource/pubsub/nats/connection_manager_test.go b/pkg/gofr/datasource/pubsub/nats/connection_manager_test.go
index 4a3e7725..e6c858e3 100644
--- a/pkg/gofr/datasource/pubsub/nats/connection_manager_test.go
+++ b/pkg/gofr/datasource/pubsub/nats/connection_manager_test.go
@@ -81,8 +81,10 @@ func TestConnectionManager_Publish(t *testing.T) {
 
 	mockJS := NewMockJetStream(ctrl)
 	mockMetrics := NewMockMetrics(ctrl)
+	mockConn := NewMockConnInterface(ctrl)
 
 	cm := &ConnectionManager{
+		conn:    mockConn,
 		jStream: mockJS,
 		logger:  logging.NewMockLogger(logging.DEBUG),
 	}
@@ -91,9 +93,12 @@ func TestConnectionManager_Publish(t *testing.T) {
 	subject := "test.subject"
 	message := []byte("test message")
 
-	mockMetrics.EXPECT().IncrementCounter(ctx, "app_pubsub_publish_total_count", "subject", subject)
-	mockJS.EXPECT().Publish(ctx, subject, message).Return(&jetstream.PubAck{}, nil)
-	mockMetrics.EXPECT().IncrementCounter(ctx, "app_pubsub_publish_success_count", "subject", subject)
+	gomock.InOrder(
+		mockMetrics.EXPECT().IncrementCounter(ctx, "app_pubsub_publish_total_count", "subject", subject),
+		mockConn.EXPECT().Status().Return(nats.CONNECTED),
+		mockJS.EXPECT().Publish(ctx, subject, message).Return(&jetstream.PubAck{}, nil),
+		mockMetrics.EXPECT().IncrementCounter(ctx, "app_pubsub_publish_success_count", "subject", subject),
+	)
 
 	err := cm.Publish(ctx, subject, message, mockMetrics)
 	require.NoError(t, err)
@@ -229,85 +234,3 @@ func TestNatsConnWrapper_NatsConn(t *testing.T) {
 
 	assert.Equal(t, mockConn, wrapper.NATSConn())
 }
-
-func TestConnectionManager_RetryConnect_Success(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	mockConn := NewMockConnInterface(ctrl)
-	mockJS := NewMockJetStream(ctrl)
-	mockNATSConnector := NewMockNATSConnector(ctrl)
-	mockJSCreator := NewMockJetStreamCreator(ctrl)
-	logger := logging.NewMockLogger(logging.DEBUG)
-
-	cm := &ConnectionManager{
-		config:           &Config{Server: "nats://localhost:4222"},
-		logger:           logger,
-		natsConnector:    mockNATSConnector,
-		jetStreamCreator: mockJSCreator,
-	}
-
-	mockNATSConnector.EXPECT().Connect(gomock.Any(), gomock.Any()).Return(mockConn, nil)
-
-	mockJSCreator.EXPECT().New(mockConn).Return(mockJS, nil)
-
-	go cm.retryConnect()
-	time.Sleep(100 * time.Millisecond) // Wait for goroutine to execute.
-
-	assert.Equal(t, mockConn, cm.conn)
-	assert.Equal(t, mockJS, cm.jStream)
-}
-
-func TestConnectionManager_RetryConnect_ConnectionFailure(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	mockNATSConnector := NewMockNATSConnector(ctrl)
-	logger := logging.NewMockLogger(logging.DEBUG)
-
-	cm := &ConnectionManager{
-		config:           &Config{Server: "nats://localhost:4222"},
-		logger:           logger,
-		natsConnector:    mockNATSConnector,
-		jetStreamCreator: NewMockJetStreamCreator(ctrl),
-	}
-
-	mockNATSConnector.EXPECT().Connect(gomock.Any(), gomock.Any()).
-		Return(nil, errConnectionError).AnyTimes()
-
-	go cm.retryConnect()
-	time.Sleep(500 * time.Millisecond) // Wait for goroutine to execute
-
-	assert.Nil(t, cm.conn)
-	assert.Nil(t, cm.jStream)
-}
-
-func TestConnectionManager_RetryConnect_JetStreamCreationFailure(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	mockConn := NewMockConnInterface(ctrl)
-	mockNATSConnector := NewMockNATSConnector(ctrl)
-	mockJSCreator := NewMockJetStreamCreator(ctrl)
-	logger := logging.NewMockLogger(logging.DEBUG)
-
-	cm := &ConnectionManager{
-		config:           &Config{Server: "nats://localhost:4222"},
-		logger:           logger,
-		natsConnector:    mockNATSConnector,
-		jetStreamCreator: mockJSCreator,
-	}
-
-	mockConn.EXPECT().Close()
-	mockNATSConnector.EXPECT().Connect(gomock.Any(), gomock.Any()).
-		Return(mockConn, nil).AnyTimes()
-
-	mockJSCreator.EXPECT().New(mockConn).
-		Return(nil, errJetStreamNotConfigured).AnyTimes()
-
-	go cm.retryConnect()
-	time.Sleep(500 * time.Millisecond)
-
-	assert.Nil(t, cm.conn)
-	assert.Nil(t, cm.jStream)
-}