update the TestCheckBalance

internal/handlers/ussd/menuhandler.go

@@ -233,28 +233,6 @@ func (h *Handlers) SaveTemporaryPin(ctx context.Context, sym string, input []byt
 	return res, nil
 }
 
-// GetVoucherList fetches the list of vouchers and formats them
-// checks whether they are stored internally before calling the API
-func (h *Handlers) GetVoucherList(ctx context.Context, sym string, input []byte) (resource.Result, error) {
-	var res resource.Result
-
-	// check if the vouchers exist internally and if not
-	// fetch from the API
-
-	// Read vouchers from the store
-	store := h.userdataStore
-	prefixdb := storage.NewSubPrefixDb(store, []byte("token_holdings"))
-
-	voucherData, err := prefixdb.Get(ctx, []byte("tokens"))
-	if err != nil {
-		return res, nil
-	}
-
-	res.Content = string(voucherData)
-
-	return res, nil
-}
-
 func (h *Handlers) ConfirmPinChange(ctx context.Context, sym string, input []byte) (resource.Result, error) {
 	var res resource.Result
 	sessionId, ok := ctx.Value("SessionId").(string)
@@ -674,6 +652,20 @@ func (h *Handlers) CheckBalance(ctx context.Context, sym string, input []byte) (
 
 			return res, nil
 		}
+
+		return res, err
+	}
+
+	if len(activeSym) == 0 {
+		logg.Printf(logging.LVL_INFO, "Using the default sym to fetch balance")
+		balance, err := h.accountService.CheckBalance(string(publicKey))
+		if err != nil {
+			return res, err
+		}
+
+		res.Content = balance
+
+		return res, nil
 	}
 
 	activeBal, err := store.ReadEntry(ctx, sessionId, utils.DATA_ACTIVE_BAL)
@@ -1057,13 +1049,13 @@ func (h *Handlers) CheckVouchers(ctx context.Context, sym string, input []byte)
 	voucherSymbolList := strings.Join(numberedSymbols, "\n")
 	voucherBalanceList := strings.Join(numberedBalances, "\n")
 
-	prefixdb := storage.NewSubPrefixDb(store, []byte("token_holdings"))
+	prefixdb := storage.NewSubPrefixDb(store, []byte("pfx"))
-	err = prefixdb.Put(ctx, []byte("tokens"), []byte(voucherSymbolList))
+	err = prefixdb.Put(ctx, []byte("sym"), []byte(voucherSymbolList))
 	if err != nil {
 		return res, nil
 	}
 
-	err = prefixdb.Put(ctx, []byte(voucherSymbolList), []byte(voucherBalanceList))
+	err = prefixdb.Put(ctx, []byte("bal"), []byte(voucherBalanceList))
 	if err != nil {
 		return res, nil
 	}
@@ -1071,6 +1063,24 @@ func (h *Handlers) CheckVouchers(ctx context.Context, sym string, input []byte)
 	return res, nil
 }
 
+// GetVoucherList fetches the list of vouchers and formats them
+func (h *Handlers) GetVoucherList(ctx context.Context, sym string, input []byte) (resource.Result, error) {
+	var res resource.Result
+
+	// Read vouchers from the store
+	store := h.userdataStore
+	prefixdb := storage.NewSubPrefixDb(store, []byte("pfx"))
+
+	voucherData, err := prefixdb.Get(ctx, []byte("sym"))
+	if err != nil {
+		return res, nil
+	}
+
+	res.Content = string(voucherData)
+
+	return res, nil
+}
+
 // ViewVoucher retrieves the token holding and balance from the subprefixDB
 func (h *Handlers) ViewVoucher(ctx context.Context, sym string, input []byte) (resource.Result, error) {
 	var res resource.Result
@@ -1082,25 +1092,25 @@ func (h *Handlers) ViewVoucher(ctx context.Context, sym string, input []byte) (r
 		return res, fmt.Errorf("missing session")
 	}
 
+	flag_incorrect_voucher, _ := h.flagManager.GetFlag("flag_incorrect_voucher")
+
 	inputStr := string(input)
 
 	if inputStr == "0" || inputStr == "99" {
+		res.FlagReset = append(res.FlagReset, flag_incorrect_voucher)
 		return res, nil
 	}
 
-	flag_incorrect_voucher, _ := h.flagManager.GetFlag("flag_incorrect_voucher")
+	prefixdb := storage.NewSubPrefixDb(store, []byte("pfx"))
-
-	// Initialize the store and prefix database
-	prefixdb := storage.NewSubPrefixDb(store, []byte("token_holdings"))
 
 	// Retrieve the voucher symbol list
-	voucherSymbolList, err := prefixdb.Get(ctx, []byte("tokens"))
+	voucherSymbolList, err := prefixdb.Get(ctx, []byte("sym"))
 	if err != nil {
 		return res, fmt.Errorf("failed to retrieve voucher symbol list: %v", err)
 	}
 
 	// Retrieve the voucher balance list
-	voucherBalanceList, err := prefixdb.Get(ctx, []byte(voucherSymbolList))
+	voucherBalanceList, err := prefixdb.Get(ctx, []byte("bal"))
 	if err != nil {
 		return res, fmt.Errorf("failed to retrieve voucher balance list: %v", err)
 	}
@@ -1134,7 +1144,7 @@ func (h *Handlers) ViewVoucher(ctx context.Context, sym string, input []byte) (r
 		}
 	}
 
-	// If a match is found, write the temporary sym , then return the symbol and balance
+	// If a match is found, write the temporary sym, then return the symbol and balance
 	if matchedSymbol != "" && matchedBalance != "" {
 		err = store.WriteEntry(ctx, sessionId, utils.DATA_TEMPORARY_SYM, []byte(matchedSymbol))
 		if err != nil {
@@ -1153,8 +1163,9 @@ func (h *Handlers) ViewVoucher(ctx context.Context, sym string, input []byte) (r
 	return res, nil
 }
 
-// SetVoucher retrieves the temporary voucher, sets it as the active voucher and
+// SetVoucher retrieves the temporary sym and balance,
-// clears the temporary voucher/sym
+// sets them as the active data and
+// clears the temporary data
 func (h *Handlers) SetVoucher(ctx context.Context, sym string, input []byte) (resource.Result, error) {
 	var res resource.Result
 	var err error

internal/handlers/ussd/menuhandler_test.go

@@ -1560,33 +1560,72 @@ func TestValidateRecipient(t *testing.T) {
 }
 
 func TestCheckBalance(t *testing.T) {
-
+	tests := []struct {
-	mockDataStore := new(mocks.MockUserDataStore)
+		name           string
-	sessionId := "session123"
+		sessionId      string
-	publicKey := "0X13242618721"
+		publicKey      string
-	balance := "0.003 CELO"
+		activeSym      string
-
+		activeBal      string
-	expectedResult := resource.Result{
+		expectedResult resource.Result
-		Content: "0.003 CELO",
+		expectError    bool
+	}{
+		{
+			name:           "User with active sym",
+			sessionId:      "session456",
+			publicKey:      "0X98765432109",
+			activeSym:      "ETH",
+			activeBal:      "1.5",
+			expectedResult: resource.Result{Content: "1.5 ETH"},
+			expectError:    false,
+		},
+		{
+			name:           "User without active sym",
+			sessionId:      "session123",
+			publicKey:      "0X13242618721",
+			activeSym:      "",
+			activeBal:      "",
+			expectedResult: resource.Result{Content: "0.003 CELO"},
+			expectError:    false,
+		},
 	}
 
-	mockCreateAccountService := new(mocks.MockAccountService)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockDataStore := new(mocks.MockUserDataStore)
+			mockAccountService := new(mocks.MockAccountService)
+			ctx := context.WithValue(context.Background(), "SessionId", tt.sessionId)
 
-	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+			h := &Handlers{
+				userdataStore:  mockDataStore,
+				accountService: mockAccountService,
+			}
 
-	h := &Handlers{
+			// Mock calls for public key
-		userdataStore:  mockDataStore,
+			mockDataStore.On("ReadEntry", ctx, tt.sessionId, utils.DATA_PUBLIC_KEY).Return([]byte(tt.publicKey), nil)
-		accountService: mockCreateAccountService,
+
-		//flagManager:    fm.parser,
+			if tt.activeSym == "" {
+				// Mock for user without active sym
+				mockDataStore.On("ReadEntry", ctx, tt.sessionId, utils.DATA_ACTIVE_SYM).Return([]byte{}, db.ErrNotFound{})
+				mockAccountService.On("CheckBalance", tt.publicKey).Return("0.003 CELO", nil)
+			} else {
+				// Mock for user with active sym
+				mockDataStore.On("ReadEntry", ctx, tt.sessionId, utils.DATA_ACTIVE_SYM).Return([]byte(tt.activeSym), nil)
+				mockDataStore.On("ReadEntry", ctx, tt.sessionId, utils.DATA_ACTIVE_BAL).Return([]byte(tt.activeBal), nil)
+			}
+
+			res, err := h.CheckBalance(ctx, "check_balance", []byte("123456"))
+
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expectedResult, res, "Result should match expected output")
+			}
+
+			mockDataStore.AssertExpectations(t)
+			mockAccountService.AssertExpectations(t)
+		})
 	}
-	//mock call operations
-	mockDataStore.On("ReadEntry", ctx, sessionId, utils.DATA_PUBLIC_KEY).Return([]byte(publicKey), nil)
-	mockCreateAccountService.On("CheckBalance", string(publicKey)).Return(balance, nil)
-
-	res, _ := h.CheckBalance(ctx, "check_balance", []byte("123456"))
-
-	assert.Equal(t, res, expectedResult, "Result should contain flag(s) that have been reset")
-
 }
 
 func TestGetProfile(t *testing.T) {