From dc782d87a8a0cf3e324fc9f3ce7b93b01a420803 Mon Sep 17 00:00:00 2001
From: alfred-mk <alfredmwaik@gmail.com>
Date: Mon, 28 Oct 2024 11:20:06 +0300
Subject: [PATCH] Updated the TestSetVoucher

---
 internal/handlers/ussd/menuhandler.go      |  9 +++---
 internal/handlers/ussd/menuhandler_test.go | 37 +++++++++++++++++++---
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/internal/handlers/ussd/menuhandler.go b/internal/handlers/ussd/menuhandler.go
index 1385261..f04b3f8 100644
--- a/internal/handlers/ussd/menuhandler.go
+++ b/internal/handlers/ussd/menuhandler.go
@@ -1197,7 +1197,7 @@ func (h *Handlers) ViewVoucher(ctx context.Context, sym string, input []byte) (r
 	}
 
 	if err := h.storeTemporaryVoucher(ctx, sessionId, metadata); err != nil {
-		return resource.Result{}, err
+		return res, err
 	}
 
 	res.FlagReset = append(res.FlagReset, flag_incorrect_voucher)
@@ -1295,15 +1295,16 @@ func (h *Handlers) SetVoucher(ctx context.Context, sym string, input []byte) (re
 	// Get temporary data
 	tempData, err := h.getTemporaryVoucherData(ctx, sessionId)
 	if err != nil {
-		return resource.Result{}, err
+		return res, err
 	}
 
 	// Set as active and clear temporary
 	if err := h.updateVoucherData(ctx, sessionId, tempData); err != nil {
-		return resource.Result{}, err
+		return res, err
 	}
 
-	return resource.Result{Content: tempData.Symbol}, nil
+	res.Content = tempData.Symbol
+	return res, nil
 }
 
 func (h *Handlers) getTemporaryVoucherData(ctx context.Context, sessionId string) (*VoucherMetadata, error) {
diff --git a/internal/handlers/ussd/menuhandler_test.go b/internal/handlers/ussd/menuhandler_test.go
index 9bcee63..b53a289 100644
--- a/internal/handlers/ussd/menuhandler_test.go
+++ b/internal/handlers/ussd/menuhandler_test.go
@@ -1816,16 +1816,43 @@ func TestSetVoucher(t *testing.T) {
 	sessionId := "session123"
 	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
 
+	// Define the temporary voucher data
 	temporarySym := []byte("tempSym")
 	temporaryBal := []byte("tempBal")
+	temporaryDecimal := []byte("2")
+	temporaryAddress := []byte("0x123456")
 
-	// Set expectations for the mock data store
+	// Define the expected active entries
+	activeEntries := map[utils.DataTyp][]byte{
+		utils.DATA_ACTIVE_SYM:     temporarySym,
+		utils.DATA_ACTIVE_BAL:     temporaryBal,
+		utils.DATA_ACTIVE_DECIMAL: temporaryDecimal,
+		utils.DATA_ACTIVE_ADDRESS: temporaryAddress,
+	}
+
+	// Define the temporary entries to be cleared
+	tempEntries := map[utils.DataTyp][]byte{
+		utils.DATA_TEMPORARY_SYM:     []byte(""),
+		utils.DATA_TEMPORARY_BAL:     []byte(""),
+		utils.DATA_TEMPORARY_DECIMAL: []byte(""),
+		utils.DATA_TEMPORARY_ADDRESS: []byte(""),
+	}
+
+	// Mocking ReadEntry calls for temporary data retrieval
 	mockDataStore.On("ReadEntry", ctx, sessionId, utils.DATA_TEMPORARY_SYM).Return(temporarySym, nil)
 	mockDataStore.On("ReadEntry", ctx, sessionId, utils.DATA_TEMPORARY_BAL).Return(temporaryBal, nil)
-	mockDataStore.On("WriteEntry", ctx, sessionId, utils.DATA_ACTIVE_SYM, temporarySym).Return(nil)
-	mockDataStore.On("WriteEntry", ctx, sessionId, utils.DATA_ACTIVE_BAL, temporaryBal).Return(nil)
-	mockDataStore.On("WriteEntry", ctx, sessionId, utils.DATA_TEMPORARY_SYM, []byte("")).Return(nil)
-	mockDataStore.On("WriteEntry", ctx, sessionId, utils.DATA_TEMPORARY_BAL, []byte("")).Return(nil)
+	mockDataStore.On("ReadEntry", ctx, sessionId, utils.DATA_TEMPORARY_DECIMAL).Return(temporaryDecimal, nil)
+	mockDataStore.On("ReadEntry", ctx, sessionId, utils.DATA_TEMPORARY_ADDRESS).Return(temporaryAddress, nil)
+
+	// Mocking WriteEntry calls for setting active data
+	for key, value := range activeEntries {
+		mockDataStore.On("WriteEntry", ctx, sessionId, key, value).Return(nil)
+	}
+
+	// Mocking WriteEntry calls for clearing temporary data
+	for key, value := range tempEntries {
+		mockDataStore.On("WriteEntry", ctx, sessionId, key, value).Return(nil)
+	}
 
 	h := &Handlers{
 		userdataStore: mockDataStore,