From a4954ab949607d8d2285e901a5427c81bdfd1a6e Mon Sep 17 00:00:00 2001
From: Henrik Hautakoski <henrik.hautakoski@gmail.com>
Date: Sun, 4 Feb 2024 16:56:59 +0100
Subject: [PATCH] api/client_test.go: Adding test for reading rollback
 messages.

Because these don't happen that often on a real ship node. Test as much as possible with unit tests.
---
 api/client_test.go | 43 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/api/client_test.go b/api/client_test.go
index 2b28b11..6fa9d8b 100644
--- a/api/client_test.go
+++ b/api/client_test.go
@@ -1,16 +1,28 @@
 package api
 
 import (
+	"bytes"
+	"io"
 	"testing"
 
 	"github.com/eosswedenorg/thalos/api/message"
+	_ "github.com/eosswedenorg/thalos/api/message/json"
 	"github.com/stretchr/testify/assert"
 )
 
-type mockReader struct{}
+type mockReader struct {
+	r io.Reader
+}
 
 func (m mockReader) Read(channel Channel) ([]byte, error) {
-	return []byte{}, nil
+	if m.r != nil {
+		b, err := io.ReadAll(m.r)
+		if err == nil && len(b) < 1 {
+			err = io.EOF
+		}
+		return b, err
+	}
+	return []byte{}, io.EOF
 }
 
 func (m mockReader) Close() error {
@@ -59,3 +71,30 @@ func TestClient_SubscribeWithNilHandler(t *testing.T) {
 
 	assert.Error(t, err)
 }
+
+func TestClient_ReadRollback(t *testing.T) {
+	called := false
+	expected := message.RollbackMessage{
+		OldBlockNum: 1000,
+		NewBlockNum: 50,
+	}
+
+	codec, err := message.GetCodec("json")
+	assert.NoError(t, err)
+
+	payload, err := codec.Encoder(expected)
+	assert.NoError(t, err)
+
+	client := NewClient(mockReader{bytes.NewReader(payload)}, codec.Decoder)
+	client.OnRollback = func(rb message.RollbackMessage) {
+		assert.Equal(t, rb, expected)
+		called = true
+	}
+
+	err = client.Subscribe(RollbackChannel)
+	assert.NoError(t, err)
+
+	client.Run()
+
+	assert.True(t, called, "Rollback callback not called when it should have been")
+}