From 64af5d50feea243e5bde352fa33d6d1d84fea9cc Mon Sep 17 00:00:00 2001
From: YashK <yash.kurlekar2017@gmail.com>
Date: Mon, 22 Apr 2024 12:13:18 +0530
Subject: [PATCH] refactor: added tests when context is cancelled for Vote()

---
 cmd/vote_test.go | 61 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/cmd/vote_test.go b/cmd/vote_test.go
index ee8424cb..9ee9c790 100644
--- a/cmd/vote_test.go
+++ b/cmd/vote_test.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"encoding/hex"
 	"errors"
 	Types "github.com/ethereum/go-ethereum/core/types"
@@ -13,6 +14,7 @@ import (
 	"razor/utils"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/ethclient"
@@ -1212,3 +1214,62 @@ func TestHandleBlock(t *testing.T) {
 		})
 	}
 }
+
+func TestVote(t *testing.T) {
+	var (
+		config                    types.Configurations
+		client                    *ethclient.Client
+		rogueData                 types.Rogue
+		account                   types.Account
+		stakerId                  uint32
+		backupNodeActionsToIgnore []string
+	)
+	type args struct {
+		header *Types.Header
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "Test when context is cancelled",
+			args: args{
+				header: &Types.Header{
+					Number: big.NewInt(101),
+				},
+			},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			SetUpMockInterfaces()
+
+			clientUtilsMock.On("GetLatestBlockWithRetry", mock.Anything).Return(tt.args.header, nil)
+			cmdUtilsMock.On("HandleBlock", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+			ut := &UtilsStruct{}
+			errChan := make(chan error)
+			// Run Vote function in a goroutine
+			go func() {
+				errChan <- ut.Vote(ctx, config, client, rogueData, account, stakerId, backupNodeActionsToIgnore)
+			}()
+
+			// Wait for some time to allow Vote function to execute
+			time.Sleep(time.Second * 2)
+
+			// Cancel the context to simulate its done
+			cancel()
+
+			// Check the error returned from the function
+			err := <-errChan
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Vote() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}