From 1278b234a56b8338afee89c250fcd25a934ccbd7 Mon Sep 17 00:00:00 2001
From: Paolo Galli <paolo.galli@vechain.org>
Date: Tue, 28 May 2024 12:16:17 +0200
Subject: [PATCH 1/3] fix: fix and re-add subscription test

---
 api/subscriptions/subscriptions_test.go | 39 +++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/api/subscriptions/subscriptions_test.go b/api/subscriptions/subscriptions_test.go
index 589a6d58d..22cabe1a8 100644
--- a/api/subscriptions/subscriptions_test.go
+++ b/api/subscriptions/subscriptions_test.go
@@ -13,6 +13,7 @@ import (
 	"net/url"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/gorilla/mux"
 	"github.com/gorilla/websocket"
@@ -34,6 +35,7 @@ func TestMain(t *testing.T) {
 	initSubscriptionsServer(t)
 	defer ts.Close()
 
+	testHandlePendingTransactions(t)
 	testHandleSubjectWithBlock(t)
 	testHandleSubjectWithEvent(t)
 	testHandleSubjectWithTransfer(t)
@@ -42,6 +44,43 @@ func TestMain(t *testing.T) {
 	testHandleSubjectWithNonValidArgument(t)
 }
 
+func testHandlePendingTransactions(t *testing.T) {
+	// This channel makes sure the new tx is notified to mempool subscribers
+	// and then to pendingTx as well so that websocket has the tx to read
+	txChan := make(chan *txpool.TxEvent)
+	sub := txPool.SubscribeTxEvent(txChan)
+	defer sub.Unsubscribe()
+
+	u := url.URL{Scheme: "ws", Host: strings.TrimPrefix(ts.URL, "http://"), Path: "/subscriptions/txpool"}
+
+	conn, resp, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	assert.NoError(t, err)
+	defer conn.Close()
+
+	// Check the protocol upgrade to websocket
+	assert.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode)
+	assert.Equal(t, "Upgrade", resp.Header.Get("Connection"))
+	assert.Equal(t, "websocket", resp.Header.Get("Upgrade"))
+
+	// Add a new tx to the mempool
+	transaction := createTx(t, repo, 1)
+	txPool.AddLocal(transaction)
+
+	<-txChan
+	time.Sleep(100 * time.Millisecond)
+
+	_, msg, err := conn.ReadMessage()
+
+	assert.NoError(t, err)
+
+	var pendingTx *PendingTxIDMessage
+	if err := json.Unmarshal(msg, &pendingTx); err != nil {
+		t.Fatal(err)
+	} else {
+		assert.Equal(t, transaction.ID(), pendingTx.ID)
+	}
+}
+
 func testHandleSubjectWithBlock(t *testing.T) {
 	genesisBlock := blocks[0]
 	queryArg := fmt.Sprintf("pos=%s", genesisBlock.Header().ID().String())

From 3a7e503b12140fb872ae4b6c2d463b40b0087bd4 Mon Sep 17 00:00:00 2001
From: Paolo Galli <paolo.galli@vechain.org>
Date: Tue, 28 May 2024 12:20:51 +0200
Subject: [PATCH 2/3] test: add comment

---
 api/subscriptions/subscriptions_test.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/api/subscriptions/subscriptions_test.go b/api/subscriptions/subscriptions_test.go
index 22cabe1a8..44cce69e2 100644
--- a/api/subscriptions/subscriptions_test.go
+++ b/api/subscriptions/subscriptions_test.go
@@ -66,6 +66,7 @@ func testHandlePendingTransactions(t *testing.T) {
 	transaction := createTx(t, repo, 1)
 	txPool.AddLocal(transaction)
 
+	// Wait for the tx to be notified from mempool
 	<-txChan
 	time.Sleep(100 * time.Millisecond)
 

From 706858b80f9f4555badeba200d74700309edf36e Mon Sep 17 00:00:00 2001
From: Paolo Galli <paolo.galli@vechain.org>
Date: Tue, 28 May 2024 15:01:04 +0200
Subject: [PATCH 3/3] test: remove redundant sleep in subscriptions test

---
 api/subscriptions/subscriptions_test.go | 2 --
 1 file changed, 2 deletions(-)

diff --git a/api/subscriptions/subscriptions_test.go b/api/subscriptions/subscriptions_test.go
index 44cce69e2..2a6506442 100644
--- a/api/subscriptions/subscriptions_test.go
+++ b/api/subscriptions/subscriptions_test.go
@@ -13,7 +13,6 @@ import (
 	"net/url"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/gorilla/mux"
 	"github.com/gorilla/websocket"
@@ -68,7 +67,6 @@ func testHandlePendingTransactions(t *testing.T) {
 
 	// Wait for the tx to be notified from mempool
 	<-txChan
-	time.Sleep(100 * time.Millisecond)
 
 	_, msg, err := conn.ReadMessage()