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") +}