From 6f14463503d8b08a3c43aa6ae28d3fead3550842 Mon Sep 17 00:00:00 2001
From: Valentin Staykov <v.staykov@razorlabs.com>
Date: Tue, 29 Oct 2024 09:56:35 +0000
Subject: [PATCH] fix: tests

---
 zk/datastream/client/stream_client_test.go | 47 +++++++++++++---------
 1 file changed, 29 insertions(+), 18 deletions(-)

diff --git a/zk/datastream/client/stream_client_test.go b/zk/datastream/client/stream_client_test.go
index 265d539b83e..be04f5856d0 100644
--- a/zk/datastream/client/stream_client_test.go
+++ b/zk/datastream/client/stream_client_test.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"context"
 	"encoding/binary"
-	"errors"
 	"fmt"
 	"net"
 	"sync"
@@ -27,7 +26,7 @@ func TestStreamClientReadHeaderEntry(t *testing.T) {
 		name           string
 		input          []byte
 		expectedResult *types.HeaderEntry
-		expectedError  error
+		expectedError  string
 	}
 	testCases := []testCase{
 		{
@@ -40,13 +39,13 @@ func TestStreamClientReadHeaderEntry(t *testing.T) {
 				TotalLength:  24,
 				TotalEntries: 64,
 			},
-			expectedError: nil,
+			expectedError: "",
 		},
 		{
 			name:           "Invalid byte array length",
 			input:          []byte{20, 21, 22, 23, 24, 20},
 			expectedResult: nil,
-			expectedError:  errors.New("failed to read header bytes reading from server: unexpected EOF"),
+			expectedError:  "read header bytes reading from server: unexpected EOF",
 		},
 	}
 
@@ -64,7 +63,11 @@ func TestStreamClientReadHeaderEntry(t *testing.T) {
 			}()
 
 			header, err := c.readHeaderEntry()
-			require.Equal(t, testCase.expectedError, err)
+			if testCase.expectedError != "" {
+				require.EqualError(t, err, testCase.expectedError)
+			} else {
+				require.NoError(t, err)
+			}
 			assert.DeepEqual(t, testCase.expectedResult, header)
 		})
 	}
@@ -75,7 +78,7 @@ func TestStreamClientReadResultEntry(t *testing.T) {
 		name           string
 		input          []byte
 		expectedResult *types.ResultEntry
-		expectedError  error
+		expectedError  string
 	}
 	testCases := []testCase{
 		{
@@ -87,7 +90,7 @@ func TestStreamClientReadResultEntry(t *testing.T) {
 				ErrorNum:   0,
 				ErrorStr:   []byte{},
 			},
-			expectedError: nil,
+			expectedError: "",
 		},
 		{
 			name:  "Happy path - error str length",
@@ -98,19 +101,19 @@ func TestStreamClientReadResultEntry(t *testing.T) {
 				ErrorNum:   0,
 				ErrorStr:   []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
 			},
-			expectedError: nil,
+			expectedError: "",
 		},
 		{
 			name:           "Invalid byte array length",
 			input:          []byte{20, 21, 22, 23, 24, 20},
 			expectedResult: nil,
-			expectedError:  errors.New("failed to read main result bytes reading from server: unexpected EOF"),
+			expectedError:  "read main result bytes reading from server: unexpected EOF",
 		},
 		{
 			name:           "Invalid error length",
 			input:          []byte{0, 0, 0, 12, 0, 0, 0, 0, 20, 21},
 			expectedResult: nil,
-			expectedError:  errors.New("failed to read result errStr bytes reading from server: unexpected EOF"),
+			expectedError:  "read result errStr bytes reading from server: unexpected EOF",
 		},
 	}
 
@@ -128,7 +131,11 @@ func TestStreamClientReadResultEntry(t *testing.T) {
 			}()
 
 			result, err := c.readResultEntry([]byte{1})
-			require.Equal(t, testCase.expectedError, err)
+			if testCase.expectedError != "" {
+				require.EqualError(t, err, testCase.expectedError)
+			} else {
+				require.NoError(t, err)
+			}
 			assert.DeepEqual(t, testCase.expectedResult, result)
 		})
 	}
@@ -139,7 +146,7 @@ func TestStreamClientReadFileEntry(t *testing.T) {
 		name           string
 		input          []byte
 		expectedResult *types.FileEntry
-		expectedError  error
+		expectedError  string
 	}
 	testCases := []testCase{
 		{
@@ -152,7 +159,7 @@ func TestStreamClientReadFileEntry(t *testing.T) {
 				EntryNum:   45,
 				Data:       []byte{0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 64},
 			},
-			expectedError: nil,
+			expectedError: "",
 		}, {
 			name:  "Happy path - no data",
 			input: []byte{2, 0, 0, 0, 17, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 45},
@@ -163,24 +170,24 @@ func TestStreamClientReadFileEntry(t *testing.T) {
 				EntryNum:   45,
 				Data:       []byte{},
 			},
-			expectedError: nil,
+			expectedError: "",
 		},
 		{
 			name:           "Invalid packet type",
 			input:          []byte{5, 0, 0, 0, 17, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 45},
 			expectedResult: nil,
-			expectedError:  errors.New("expected data packet type 2 or 254 and received 5"),
+			expectedError:  "expected data packet type 2 or 254 and received 5",
 		},
 		{
 			name:           "Invalid byte array length",
 			input:          []byte{2, 21, 22, 23, 24, 20},
 			expectedResult: nil,
-			expectedError:  errors.New("error reading file bytes: reading from server: unexpected EOF"),
+			expectedError:  "reading file bytes: reading from server: unexpected EOF",
 		}, {
 			name:           "Invalid data length",
 			input:          []byte{2, 0, 0, 0, 31, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 45, 0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 64},
 			expectedResult: nil,
-			expectedError:  errors.New("error reading file data bytes: reading from server: unexpected EOF"),
+			expectedError:  "reading file data bytes: reading from server: unexpected EOF",
 		},
 	}
 	for _, testCase := range testCases {
@@ -197,7 +204,11 @@ func TestStreamClientReadFileEntry(t *testing.T) {
 			}()
 
 			result, err := c.NextFileEntry()
-			require.Equal(t, testCase.expectedError, err)
+			if testCase.expectedError != "" {
+				require.EqualError(t, err, testCase.expectedError)
+			} else {
+				require.NoError(t, err)
+			}
 			assert.DeepEqual(t, testCase.expectedResult, result)
 		})
 	}