cleanup

cmd/main.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"log"
 	"os"
 	"path"
 

internal/handlers/ussd/menuhandler.go

@@ -61,13 +61,15 @@ type Handlers struct {
 	pe             *persist.Persister
 	st             *state.State
 	ca             cache.Memory
-	userdataStore      db.Db
+	userdataStore  utils.DataStore
 	flagManager    *asm.FlagParser
-	accountFileHandler *utils.AccountFileHandler
 	accountService server.AccountServiceInterface
 }
 
-func NewHandlers(appFlags *asm.FlagParser, pe *persist.Persister, userdataStore db.Db) (*Handlers, error) {
+func NewHandlers(parser *asm.FlagParser, pe *persist.Persister, userdataStore db.Db) (*Handlers, error) {
+	userDb := utils.UserDataStore{
+		Store: userdataStore,
+	}
 	if pe == nil {
 		return nil, fmt.Errorf("cannot create handler with nil persister")
 	}
@@ -76,9 +78,8 @@ func NewHandlers(appFlags *asm.FlagParser, pe *persist.Persister, userdataStore
 	}
 	h := &Handlers{
 		pe:             pe,
-		userdataStore:      userdataStore,
-		flagManager:        appFlags,
-		accountFileHandler: utils.NewAccountFileHandler(userdataStore),
+		userdataStore:  userDb,
+		flagManager:    parser,
 		accountService: &server.AccountService{},
 	}
 	return h, nil
@@ -145,8 +146,10 @@ func (h *Handlers) createAccountNoExist(ctx context.Context, sessionId string, r
 		utils.DATA_PUBLIC_KEY:   accountResp.Result.PublicKey,
 		utils.DATA_CUSTODIAL_ID: accountResp.Result.CustodialId.String(),
 	}
+
 	for key, value := range data {
-		err := utils.WriteEntry(ctx, h.userdataStore, sessionId, key, []byte(value))
+		store := h.userdataStore
+		err := store.WriteEntry(ctx, sessionId, key, []byte(value))
 		if err != nil {
 			return err
 		}
@@ -167,7 +170,8 @@ func (h *Handlers) CreateAccount(ctx context.Context, sym string, input []byte)
 	if !ok {
 		return res, fmt.Errorf("missing session")
 	}
-	_, err = utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_ACCOUNT_CREATED)
+	store := h.userdataStore
+	_, err = store.ReadEntry(ctx, sessionId, utils.DATA_ACCOUNT_CREATED)
 	if err != nil {
 		if db.IsNotFound(err) {
 			logg.Printf(logging.LVL_INFO, "Creating an account because it doesn't exist")
@@ -200,12 +204,11 @@ func (h *Handlers) SavePin(ctx context.Context, sym string, input []byte) (resou
 	}
 
 	res.FlagReset = append(res.FlagReset, flag_incorrect_pin)
-
-	err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_ACCOUNT_PIN, []byte(accountPIN))
+	store := h.userdataStore
+	err = store.WriteEntry(ctx, sessionId, utils.DATA_ACCOUNT_PIN, []byte(accountPIN))
 	if err != nil {
 		return res, err
 	}
-
 	return res, nil
 }
 
@@ -250,7 +253,12 @@ func (h *Handlers) VerifyPin(ctx context.Context, sym string, input []byte) (res
 		return res, fmt.Errorf("missing session")
 	}
 
-	AccountPin, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_ACCOUNT_PIN)
+	//AccountPin, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_ACCOUNT_PIN)
+	store := h.userdataStore.(utils.UserDataStore)
+	AccountPin, err := store.ReadEntry(ctx, sessionId, utils.DATA_ACCOUNT_PIN)
+	if err != nil {
+		return res, err
+	}
 
 	if bytes.Equal(input, AccountPin) {
 		res.FlagSet = []uint32{flag_valid_pin}
@@ -284,9 +292,10 @@ func (h *Handlers) SaveFirstname(ctx context.Context, sym string, input []byte)
 
 	if len(input) > 0 {
 		firstName := string(input)
-		err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_FIRST_NAME, []byte(firstName))
+		store := h.userdataStore
+		err = store.WriteEntry(ctx, sessionId, utils.DATA_FIRST_NAME, []byte(firstName))
 		if err != nil {
-			return res, nil
+			return res, err
 		}
 	}
 
@@ -304,7 +313,11 @@ func (h *Handlers) SaveFamilyname(ctx context.Context, sym string, input []byte)
 
 	if len(input) > 0 {
 		familyName := string(input)
-		err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_FAMILY_NAME, []byte(familyName))
+		store := h.userdataStore
+		err = store.WriteEntry(ctx, sessionId, utils.DATA_FAMILY_NAME, []byte(familyName))
+		if err != nil {
+			return res, err
+		}
 		if err != nil {
 			return res, nil
 		}
@@ -326,9 +339,10 @@ func (h *Handlers) SaveYob(ctx context.Context, sym string, input []byte) (resou
 
 	if len(input) == 4 {
 		yob := string(input)
-		err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_YOB, []byte(yob))
+		store := h.userdataStore
+		err = store.WriteEntry(ctx, sessionId, utils.DATA_YOB, []byte(yob))
 		if err != nil {
-			return res, nil
+			return res, err
 		}
 	}
 
@@ -346,9 +360,10 @@ func (h *Handlers) SaveLocation(ctx context.Context, sym string, input []byte) (
 
 	if len(input) > 0 {
 		location := string(input)
-		err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_LOCATION, []byte(location))
+		store := h.userdataStore
+		err = store.WriteEntry(ctx, sessionId, utils.DATA_LOCATION, []byte(location))
 		if err != nil {
-			return res, nil
+			return res, err
 		}
 	}
 
@@ -374,7 +389,8 @@ func (h *Handlers) SaveGender(ctx context.Context, sym string, input []byte) (re
 		case "3":
 			gender = "Unspecified"
 		}
-		err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_GENDER, []byte(gender))
+		store := h.userdataStore
+		err = store.WriteEntry(ctx, sessionId, utils.DATA_GENDER, []byte(gender))
 		if err != nil {
 			return res, nil
 		}
@@ -394,7 +410,8 @@ func (h *Handlers) SaveOfferings(ctx context.Context, sym string, input []byte)
 
 	if len(input) > 0 {
 		offerings := string(input)
-		err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_OFFERINGS, []byte(offerings))
+		store := h.userdataStore
+		err = store.WriteEntry(ctx, sessionId, utils.DATA_OFFERINGS, []byte(offerings))
 		if err != nil {
 			return res, nil
 		}
@@ -432,7 +449,8 @@ func (h *Handlers) CheckIdentifier(ctx context.Context, sym string, input []byte
 		return res, fmt.Errorf("missing session")
 	}
 
-	publicKey, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_PUBLIC_KEY)
+	store := h.userdataStore
+	publicKey, _ := store.ReadEntry(ctx, sessionId, utils.DATA_PUBLIC_KEY)
 
 	res.Content = string(publicKey)
 
@@ -454,7 +472,11 @@ func (h *Handlers) Authorize(ctx context.Context, sym string, input []byte) (res
 	flag_account_authorized, _ := h.flagManager.GetFlag("flag_account_authorized")
 	flag_allow_update, _ := h.flagManager.GetFlag("flag_allow_update")
 
-	AccountPin, err := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_ACCOUNT_PIN)
+	store := h.userdataStore
+	AccountPin, err := store.ReadEntry(ctx, sessionId, utils.DATA_ACCOUNT_PIN)
+	if err != nil {
+		return res, err
+	}
 
 	if err == nil {
 		if len(input) == 4 {
@@ -501,8 +523,11 @@ func (h *Handlers) CheckAccountStatus(ctx context.Context, sym string, input []b
 	if !ok {
 		return res, fmt.Errorf("missing session")
 	}
-
-	trackingId, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_TRACKING_ID)
+	store := h.userdataStore.(utils.UserDataStore)
+	trackingId, err := store.ReadEntry(ctx, sessionId, utils.DATA_TRACKING_ID)
+	if err != nil {
+		return res, err
+	}
 
 	status, err := h.accountService.CheckAccountStatus(string(trackingId))
 	if err != nil {
@@ -510,7 +535,7 @@ func (h *Handlers) CheckAccountStatus(ctx context.Context, sym string, input []b
 		return res, err
 	}
 
-	err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_ACCOUNT_STATUS, []byte(status))
+	err = store.WriteEntry(ctx, sessionId, utils.DATA_ACCOUNT_STATUS, []byte(status))
 	if err != nil {
 		return res, nil
 	}
@@ -585,7 +610,11 @@ func (h *Handlers) CheckBalance(ctx context.Context, sym string, input []byte) (
 		return res, fmt.Errorf("missing session")
 	}
 
-	publicKey, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_PUBLIC_KEY)
+	store := h.userdataStore.(utils.UserDataStore)
+	publicKey, err := store.ReadEntry(ctx, sessionId, utils.DATA_PUBLIC_KEY)
+	if err != nil {
+		return res, err
+	}
 
 	balance, err := h.accountService.CheckBalance(string(publicKey))
 	if err != nil {
@@ -618,8 +647,8 @@ func (h *Handlers) ValidateRecipient(ctx context.Context, sym string, input []by
 
 			return res, nil
 		}
-
-		err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_RECIPIENT, []byte(recipient))
+		store := h.userdataStore
+		err = store.WriteEntry(ctx, sessionId, utils.DATA_RECIPIENT, []byte(recipient))
 		if err != nil {
 			return res, nil
 		}
@@ -641,13 +670,13 @@ func (h *Handlers) TransactionReset(ctx context.Context, sym string, input []byt
 
 	flag_invalid_recipient, _ := h.flagManager.GetFlag("flag_invalid_recipient")
 	flag_invalid_recipient_with_invite, _ := h.flagManager.GetFlag("flag_invalid_recipient_with_invite")
-
-	err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_AMOUNT, []byte(""))
+	store := h.userdataStore
+	err = store.WriteEntry(ctx, sessionId, utils.DATA_AMOUNT, []byte(""))
 	if err != nil {
 		return res, nil
 	}
 
-	err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_RECIPIENT, []byte(""))
+	err = store.WriteEntry(ctx, sessionId, utils.DATA_RECIPIENT, []byte(""))
 	if err != nil {
 		return res, nil
 	}
@@ -668,8 +697,8 @@ func (h *Handlers) ResetTransactionAmount(ctx context.Context, sym string, input
 	}
 
 	flag_invalid_amount, _ := h.flagManager.GetFlag("flag_invalid_amount")
-
-	err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_AMOUNT, []byte(""))
+	store := h.userdataStore
+	err = store.WriteEntry(ctx, sessionId, utils.DATA_AMOUNT, []byte(""))
 	if err != nil {
 		return res, nil
 	}
@@ -689,8 +718,8 @@ func (h *Handlers) MaxAmount(ctx context.Context, sym string, input []byte) (res
 	if !ok {
 		return res, fmt.Errorf("missing session")
 	}
-
-	publicKey, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_PUBLIC_KEY)
+	store := h.userdataStore
+	publicKey, _ := store.ReadEntry(ctx, sessionId, utils.DATA_PUBLIC_KEY)
 
 	balance, err := h.accountService.CheckBalance(string(publicKey))
 	if err != nil {
@@ -759,8 +788,8 @@ func (h *Handlers) ValidateAmount(ctx context.Context, sym string, input []byte)
 	}
 
 	res.Content = fmt.Sprintf("%.3f", inputAmount) // Format to 3 decimal places
-
-	err = utils.WriteEntry(ctx, h.userdataStore, sessionId, utils.DATA_AMOUNT, []byte(amountStr))
+	store := h.userdataStore
+	err = store.WriteEntry(ctx, sessionId, utils.DATA_AMOUNT, []byte(amountStr))
 	if err != nil {
 		return res, err
 	}
@@ -776,8 +805,8 @@ func (h *Handlers) GetRecipient(ctx context.Context, sym string, input []byte) (
 	if !ok {
 		return res, fmt.Errorf("missing session")
 	}
-
-	recipient, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_RECIPIENT)
+	store := h.userdataStore
+	recipient, _ := store.ReadEntry(ctx, sessionId, utils.DATA_RECIPIENT)
 
 	res.Content = string(recipient)
 
@@ -793,7 +822,8 @@ func (h *Handlers) GetSender(ctx context.Context, sym string, input []byte) (res
 		return res, fmt.Errorf("missing session")
 	}
 
-	publicKey, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_PUBLIC_KEY)
+	store := h.userdataStore
+	publicKey, _ := store.ReadEntry(ctx, sessionId, utils.DATA_PUBLIC_KEY)
 
 	res.Content = string(publicKey)
 
@@ -808,8 +838,8 @@ func (h *Handlers) GetAmount(ctx context.Context, sym string, input []byte) (res
 	if !ok {
 		return res, fmt.Errorf("missing session")
 	}
-
-	amount, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_AMOUNT)
+	store := h.userdataStore
+	amount, _ := store.ReadEntry(ctx, sessionId, utils.DATA_AMOUNT)
 
 	res.Content = string(amount)
 
@@ -832,8 +862,11 @@ func (h *Handlers) QuitWithBalance(ctx context.Context, sym string, input []byte
 	l := gotext.NewLocale(translationDir, code)
 	l.AddDomain("default")
 
-	publicKey, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_PUBLIC_KEY)
-
+	store := h.userdataStore.(utils.UserDataStore)
+	publicKey, err := store.ReadEntry(ctx, sessionId, utils.DATA_PUBLIC_KEY)
+	if err != nil {
+		return res, err
+	}
 	balance, err := h.accountService.CheckBalance(string(publicKey))
 	if err != nil {
 		return res, nil
@@ -858,12 +891,12 @@ func (h *Handlers) InitiateTransaction(ctx context.Context, sym string, input []
 	l.AddDomain("default")
 	// TODO
 	// Use the amount, recipient and sender to call the API and initialize the transaction
+	store := h.userdataStore
+	publicKey, _ := store.ReadEntry(ctx, sessionId, utils.DATA_PUBLIC_KEY)
 
-	publicKey, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_PUBLIC_KEY)
+	amount, _ := store.ReadEntry(ctx, sessionId, utils.DATA_AMOUNT)
 
-	amount, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_AMOUNT)
-
-	recipient, _ := utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_RECIPIENT)
+	recipient, _ := store.ReadEntry(ctx, sessionId, utils.DATA_RECIPIENT)
 
 	res.Content = l.Get("Your request has been sent. %s will receive %s from %s.", string(recipient), string(amount), string(publicKey))
 
@@ -894,14 +927,14 @@ func (h *Handlers) GetProfileInfo(ctx context.Context, sym string, input []byte)
 		}
 		return string(entry)
 	}
-
+	store := h.userdataStore
 	// Retrieve user data as strings with fallback to defaultValue
-	firstName := getEntryOrDefault(utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_FIRST_NAME))
-	familyName := getEntryOrDefault(utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_FAMILY_NAME))
-	yob := getEntryOrDefault(utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_YOB))
-	gender := getEntryOrDefault(utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_GENDER))
-	location := getEntryOrDefault(utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_LOCATION))
-	offerings := getEntryOrDefault(utils.ReadEntry(ctx, h.userdataStore, sessionId, utils.DATA_OFFERINGS))
+	firstName := getEntryOrDefault(store.ReadEntry(ctx, sessionId, utils.DATA_FIRST_NAME))
+	familyName := getEntryOrDefault(store.ReadEntry(ctx, sessionId, utils.DATA_FAMILY_NAME))
+	yob := getEntryOrDefault(store.ReadEntry(ctx, sessionId, utils.DATA_YOB))
+	gender := getEntryOrDefault(store.ReadEntry(ctx, sessionId, utils.DATA_GENDER))
+	location := getEntryOrDefault(store.ReadEntry(ctx, sessionId, utils.DATA_LOCATION))
+	offerings := getEntryOrDefault(store.ReadEntry(ctx, sessionId, utils.DATA_OFFERINGS))
 
 	// Construct the full name
 	name := defaultValue

internal/handlers/ussd/menuhandler_test.go

@@ -2,170 +2,374 @@ package ussd
 
 import (
 	"context"
-	"encoding/json"
 	"testing"
 
-	"git.defalsify.org/vise.git/db"
+	"git.defalsify.org/vise.git/resource"
 	"git.grassecon.net/urdt/ussd/internal/handlers/ussd/mocks"
-	"git.grassecon.net/urdt/ussd/internal/models"
 	"git.grassecon.net/urdt/ussd/internal/utils"
 	"github.com/alecthomas/assert/v2"
-	"github.com/stretchr/testify/mock"
 )
 
-func TestCreateAccount_Success(t *testing.T) {
-	mockCreateAccountService := new(mocks.MockAccountService)
-	mockUserDataStore := new(mocks.MockDb)
+// func TestCreateAccount(t *testing.T) {
+// 	// Create a new instance of MockMyDataStore
+// 	mockDataStore := new(mocks.MockMyDataStore)
+// 	mockCreateAccountService := new(mocks.MockAccountService)
 
+// 	// Define session ID and mock data
+// 	sessionId := "session123"
+// 	typ := utils.DATA_ACCOUNT_CREATED
+// 	fakeError := db.ErrNotFound{}
+// 	// Create context with session ID
+// 	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+
+// 	// Define expected interactions with the mock
+// 	mockDataStore.On("ReadEntry", ctx, sessionId, typ).Return([]byte("123"), fakeError)
+// 	expectedAccountResp := &models.AccountResponse{
+// 		Ok: true,
+// 		Result: struct {
+// 			CustodialId json.Number `json:"custodialId"`
+// 			PublicKey   string      `json:"publicKey"`
+// 			TrackingId  string      `json:"trackingId"`
+// 		}{
+// 			CustodialId: "12",
+// 			PublicKey:   "0x8E0XSCSVA",
+// 			TrackingId:  "d95a7e83-196c-4fd0-866fSGAGA",
+// 		},
+// 	}
+// 	mockCreateAccountService.On("CreateAccount").Return(expectedAccountResp, nil)
+// 	data := map[utils.DataTyp]string{
+// 		utils.DATA_TRACKING_ID:  expectedAccountResp.Result.TrackingId,
+// 		utils.DATA_PUBLIC_KEY:   expectedAccountResp.Result.PublicKey,
+// 		utils.DATA_CUSTODIAL_ID: expectedAccountResp.Result.CustodialId.String(),
+// 	}
+
+// 	for key, value := range data {
+// 		//err := utils.WriteEntry(ctx, h.userdataStore, sessionId, key, []byte(value))
+// 		mockDataStore.On("WriteEntry", ctx, sessionId, key, []byte(value)).Return(nil)
+// 	}
+// 	//mockDataStore.On("WriteEntry", mock.Anything, sessionId, mock.Anything, mock.Anything).Return(nil)
+
+// 	// Create a Handlers instance with the mock data store
+// 	h := &Handlers{
+// 		userdataStore:  mockDataStore,
+// 		accountService: mockCreateAccountService,
+// 	}
+
+// 	// Call the method you want to test
+// 	_, err := h.CreateAccount(ctx, "some-symbol", []byte("some-input"))
+
+// 	// Assert that no errors occurred
+// 	assert.NoError(t, err)
+
+// 	// Assert that expectations were met
+// 	mockDataStore.AssertExpectations(t)
+// }
+
+func TestSaveFirstname(t *testing.T) {
+	// Create a new instance of MockMyDataStore
+	mockStore := new(mocks.MockUserDataStore)
+
+	// Define test data
+	sessionId := "session123"
+	firstName := "John"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+
+	// Set up the expected behavior of the mock
+	mockStore.On("WriteEntry", ctx, sessionId, utils.DATA_FIRST_NAME, []byte(firstName)).Return(nil)
+
+	// Create the Handlers instance with the mock store
 	h := &Handlers{
-		userdataStore:  mockUserDataStore,
-		accountService: mockCreateAccountService,
+		userdataStore: mockStore,
 	}
-	ctx := context.WithValue(context.Background(), "SessionId", "test-session-12345")
-	k := utils.PackKey(utils.DATA_ACCOUNT_CREATED, []byte("test-session-12345"))
-	mockUserDataStore.On("SetPrefix", uint8(0x20)).Return(nil)
-	mockUserDataStore.On("SetSession", "test-session-12345").Return(nil)
-	mockUserDataStore.On("Get", ctx, k).
-		Return(nil, db.ErrNotFound{})
 
-	// Define expected account response after api call
-	expectedAccountResp := &models.AccountResponse{
-		Ok: true,
-		Result: struct {
-			CustodialId json.Number `json:"custodialId"`
-			PublicKey   string      `json:"publicKey"`
-			TrackingId  string      `json:"trackingId"`
-		}{
-			CustodialId: "12",
-			PublicKey:   "0x8E0XSCSVA",
-			TrackingId:  "d95a7e83-196c-4fd0-866fSGAGA",
-		},
-	}
-	mockCreateAccountService.On("CreateAccount").Return(expectedAccountResp, nil)
-
-	_, err := h.CreateAccount(ctx, "create_account", []byte("create_account"))
+	// Call the method
+	res, err := h.SaveFirstname(ctx, "save_firstname", []byte(firstName))
 
 	// Assert results
 	assert.NoError(t, err)
-	assert.Equal(t, expectedAccountResp.Ok, true)
-}
+	assert.Equal(t, resource.Result{}, res)
 
-func TestSaveFirstname(t *testing.T) {
-	// Create a mock database
-	mockDb := new(mocks.MockDb)
-
-	// Create a Handlers instance with the mock database
-	h := &Handlers{
-		userdataStore: mockDb,
-	}
-
-	// Create a context with a session ID
-	ctx := context.WithValue(context.Background(), "SessionId", "test-session")
-
-	tests := []struct {
-		name        string
-		input       []byte
-		expectError bool
-		setupMock   func(*mocks.MockDb)
-	}{
-		{
-			name:        "Valid first name",
-			input:       []byte("John"),
-			expectError: false,
-			setupMock: func(m *mocks.MockDb) {
-				m.On("SetPrefix", uint8(0x20)).Return(nil)
-				m.On("SetSession", "test-session").Return(nil)
-				m.On("Put", mock.Anything, mock.Anything, []byte("John")).Return(nil)
-			},
-		},
-		{
-			name:        "Empty first name",
-			input:       []byte{},
-			expectError: false, // Note: The function doesn't return an error for empty input
-			setupMock:   func(m *mocks.MockDb) {},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			// Setup mock expectations
-			tt.setupMock(mockDb)
-
-			// Call the function
-			_, err := h.SaveFirstname(ctx, "", tt.input)
-
-			if tt.expectError {
-				assert.Error(t, err)
-			} else {
-				assert.NoError(t, err)
-				mockDb.AssertExpectations(t)
-			}
-
-			// Clear mock for the next test
-			mockDb.ExpectedCalls = nil
-			mockDb.Calls = nil
-		})
-	}
+	// Assert all expectations were met
+	mockStore.AssertExpectations(t)
 }
 
 func TestSaveFamilyname(t *testing.T) {
-	// Create a mock database
-	mockDb := new(mocks.MockDb)
+	// Create a new instance of MockMyDataStore
+	mockStore := new(mocks.MockUserDataStore)
 
-	// Create a Handlers instance with the mock database
+	// Define test data
+	sessionId := "session123"
+	familyName := "Doeee"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+
+	// Set up the expected behavior of the mock
+	mockStore.On("WriteEntry", ctx, sessionId, utils.DATA_FAMILY_NAME, []byte(familyName)).Return(nil)
+
+	// Create the Handlers instance with the mock store
 	h := &Handlers{
-		userdataStore: mockDb,
+		userdataStore: mockStore,
 	}
 
-	// Create a context with a session ID
-	ctx := context.WithValue(context.Background(), "SessionId", "test-session")
+	// Call the method
+	res, err := h.SaveFamilyname(ctx, "save_familyname", []byte(familyName))
 
+	// Assert results
+	assert.NoError(t, err)
+	assert.Equal(t, resource.Result{}, res)
+
+	// Assert all expectations were met
+	mockStore.AssertExpectations(t)
+}
+
+func TestSaveYoB(t *testing.T) {
+	// Create a new instance of MockMyDataStore
+	mockStore := new(mocks.MockUserDataStore)
+
+	// Define test data
+	sessionId := "session123"
+	yob := "1980"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+
+	// Set up the expected behavior of the mock
+	mockStore.On("WriteEntry", ctx, sessionId, utils.DATA_YOB, []byte(yob)).Return(nil)
+
+	// Create the Handlers instance with the mock store
+	h := &Handlers{
+		userdataStore: mockStore,
+	}
+
+	// Call the method
+	res, err := h.SaveYob(ctx, "save_yob", []byte(yob))
+
+	// Assert results
+	assert.NoError(t, err)
+	assert.Equal(t, resource.Result{}, res)
+
+	// Assert all expectations were met
+	mockStore.AssertExpectations(t)
+}
+
+func TestSaveLocation(t *testing.T) {
+	// Create a new instance of MockMyDataStore
+	mockStore := new(mocks.MockUserDataStore)
+
+	// Define test data
+	sessionId := "session123"
+	yob := "Kilifi"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+
+	// Set up the expected behavior of the mock
+	mockStore.On("WriteEntry", ctx, sessionId, utils.DATA_LOCATION, []byte(yob)).Return(nil)
+
+	// Create the Handlers instance with the mock store
+	h := &Handlers{
+		userdataStore: mockStore,
+	}
+
+	// Call the method
+	res, err := h.SaveLocation(ctx, "save_location", []byte(yob))
+
+	// Assert results
+	assert.NoError(t, err)
+	assert.Equal(t, resource.Result{}, res)
+
+	// Assert all expectations were met
+	mockStore.AssertExpectations(t)
+}
+
+func TestSaveGender(t *testing.T) {
+	// Create a new instance of MockMyDataStore
+	mockStore := new(mocks.MockUserDataStore)
+
+	// Define the session ID and context
+	sessionId := "session123"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+
+	// Define test cases
 	tests := []struct {
 		name           string
 		input          []byte
-		expectError bool
-		setupMock   func(*mocks.MockDb)
+		expectedGender string
+		expectCall     bool
 	}{
 		{
-			name:        "Valid family name",
-			input:       []byte("Smith"),
-			expectError: false,
-			setupMock: func(m *mocks.MockDb) {
-				m.On("SetPrefix", uint8(0x20)).Return(nil)
-				m.On("SetSession", "test-session").Return(nil)
-				m.On("Put", mock.Anything, mock.Anything, []byte("Smith")).Return(nil)
-			},
+			name:           "Valid Male Input",
+			input:          []byte("1"),
+			expectedGender: "Male",
+			expectCall:     true,
 		},
 		{
-			name:        "Empty family name",
-			input:       []byte{},
-			expectError: true,
-			setupMock:   func(m *mocks.MockDb) {},
+			name:           "Valid Female Input",
+			input:          []byte("2"),
+			expectedGender: "Female",
+			expectCall:     true,
+		},
+		{
+			name:           "Valid Unspecified Input",
+			input:          []byte("3"),
+			expectedGender: "Unspecified",
+			expectCall:     true,
+		},
+		{
+			name:           "Empty Input",
+			input:          []byte(""),
+			expectedGender: "",
+			expectCall:     false,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			// Setup mock expectations
-			tt.setupMock(mockDb)
-
-			// Call the function
-			_, err := h.SaveFamilyname(ctx, "", tt.input)
-
-			if tt.expectError {
-				assert.Error(t, err)
+			// Set up expectations for the mock database
+			if tt.expectCall {
+				expectedKey := utils.DATA_GENDER
+				mockStore.On("WriteEntry", ctx, sessionId, expectedKey, []byte(tt.expectedGender)).Return(nil)
 			} else {
-				assert.NoError(t, err)
-				mockDb.AssertExpectations(t)
+				mockStore.On("WriteEntry", ctx, sessionId, utils.DATA_GENDER, []byte(tt.expectedGender)).Return(nil)
 			}
 
-			// Clear mock for the next test
-			mockDb.ExpectedCalls = nil
-			mockDb.Calls = nil
+			// Create the Handlers instance with the mock store
+			h := &Handlers{
+				userdataStore: mockStore,
+			}
+
+			// Call the method
+			_, err := h.SaveGender(ctx, "someSym", tt.input)
+
+			// Assert no error
+			assert.NoError(t, err)
+
+			// Verify expectations
+			if tt.expectCall {
+				mockStore.AssertCalled(t, "WriteEntry", ctx, sessionId, utils.DATA_GENDER, []byte(tt.expectedGender))
+			} else {
+				mockStore.AssertNotCalled(t, "WriteEntry", ctx, sessionId, utils.DATA_GENDER, []byte(tt.expectedGender))
+			}
 		})
 	}
 }
 
-func testSavePin () {
+func TestCheckIdentifier(t *testing.T) {
+	// Create a new instance of MockMyDataStore
+	mockStore := new(mocks.MockUserDataStore)
+
+	// Define the session ID and context
+	sessionId := "session123"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+
+	// Define test cases
+	tests := []struct {
+		name            string
+		mockPublicKey   []byte
+		mockErr         error
+		expectedContent string
+		expectError     bool
+	}{
+		{
+			name:            "Saved public Key",
+			mockPublicKey:   []byte("0xa8363"),
+			mockErr:         nil,
+			expectedContent: "0xa8363",
+			expectError:     false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Set up expectations for the mock database
+			mockStore.On("ReadEntry", ctx, sessionId, utils.DATA_PUBLIC_KEY).Return(tt.mockPublicKey, tt.mockErr)
+
+			// Create the Handlers instance with the mock store
+			h := &Handlers{
+				userdataStore: mockStore,
+			}
+
+			// Call the method
+			res, err := h.CheckIdentifier(ctx, "check_identifier", nil)
+
+			// Assert results
+
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expectedContent, res.Content)
+
+			// Verify expectations
+			mockStore.AssertExpectations(t)
+		})
+	}
+}
+
+func TestMaxAmount(t *testing.T) {
+	mockStore := new(mocks.MockUserDataStore)
+	mockCreateAccountService := new(mocks.MockAccountService)
+
+	// Define test data
+	sessionId := "session123"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+	publicKey := "0xcasgatweksalw1018221"
+	expectedBalance := "0.003CELO"
+
+	// Set up the expected behavior of the mock
+	mockStore.On("ReadEntry", ctx, sessionId, utils.DATA_PUBLIC_KEY).Return([]byte(publicKey), nil)
+	mockCreateAccountService.On("CheckBalance", publicKey).Return(expectedBalance, nil)
+
+	// Create the Handlers instance with the mock store
+	h := &Handlers{
+		userdataStore:  mockStore,
+		accountService: mockCreateAccountService,
+	}
+
+	// Call the method
+	res, _ := h.MaxAmount(ctx, "max_amount", []byte("check_balance"))
+
+	//Assert that the balance that was set as the result content is what was returned by  Check Balance
+	assert.Equal(t, expectedBalance, res.Content)
+
+}
+
+func TestGetSender(t *testing.T) {
+	mockStore := new(mocks.MockUserDataStore)
+
+	// Define test data
+	sessionId := "session123"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+	publicKey := "0xcasgatweksalw1018221"
+
+	// Set up the expected behavior of the mock
+	mockStore.On("ReadEntry", ctx, sessionId, utils.DATA_PUBLIC_KEY).Return([]byte(publicKey), nil)
+
+	// Create the Handlers instance with the mock store
+	h := &Handlers{
+		userdataStore: mockStore,
+	}
+
+	// Call the method
+	res, _ := h.GetSender(ctx, "max_amount", []byte("check_balance"))
+
+	//Assert that the public key from readentry operation  is what was set as the result content.
+	assert.Equal(t, publicKey, res.Content)
+
+}
+
+func TestGetAmount(t *testing.T) {
+	mockStore := new(mocks.MockUserDataStore)
+
+	// Define test data
+	sessionId := "session123"
+	ctx := context.WithValue(context.Background(), "SessionId", sessionId)
+	Amount := "0.03CELO"
+
+	// Set up the expected behavior of the mock
+	mockStore.On("ReadEntry", ctx, sessionId, utils.DATA_AMOUNT).Return([]byte(Amount), nil)
+
+	// Create the Handlers instance with the mock store
+	h := &Handlers{
+		userdataStore: mockStore,
+	}
+
+	// Call the method
+	res, _ := h.GetAmount(ctx, "get_amount", []byte("Getting amount..."))
+
+	//Assert that the retrieved amount is what was set as the content
+	assert.Equal(t, Amount, res.Content)
 
 }

internal/handlers/ussd/mocks/userdbmock.go

@@ -0,0 +1,69 @@
+package mocks
+
+import (
+	"context"
+
+	"git.defalsify.org/vise.git/lang"
+	"git.grassecon.net/urdt/ussd/internal/utils"
+	"github.com/stretchr/testify/mock"
+)
+
+type MockUserDataStore struct {
+	mock.Mock
+}
+
+func (m *MockUserDataStore) SetPrefix(prefix uint8) {
+	m.Called(prefix)
+}
+
+func (m *MockUserDataStore) SetSession(sessionId string) {
+	m.Called(sessionId)
+}
+
+func (m *MockUserDataStore) Get(ctx context.Context, key []byte) ([]byte, error) {
+	args := m.Called(ctx, key)
+	return args.Get(0).([]byte), args.Error(1)
+}
+
+func (m *MockUserDataStore) ReadEntry(ctx context.Context, sessionId string, typ utils.DataTyp) ([]byte, error) {
+	args := m.Called(ctx, sessionId, typ)
+	return args.Get(0).([]byte), args.Error(1)
+}
+
+func (m *MockUserDataStore) WriteEntry(ctx context.Context, sessionId string, typ utils.DataTyp, value []byte) error {
+	args := m.Called(ctx, sessionId, typ, value)
+	return args.Error(0)
+}
+
+func (m *MockUserDataStore) Prefix() uint8 {
+	args := m.Called()
+	return args.Get(0).(uint8)
+}
+
+func (m *MockUserDataStore) Safe() bool {
+	args := m.Called()
+	return args.Get(0).(bool)
+}
+
+func (m *MockUserDataStore) SetLanguage(language *lang.Language) {
+	m.Called(language)
+}
+
+func (m *MockUserDataStore) SetLock(uint8, bool) error {
+	args := m.Called()
+	return args.Error(0)
+}
+
+func (m *MockUserDataStore) Connect(ctx context.Context, connectionStr string) error {
+	args := m.Called(ctx, connectionStr)
+	return args.Error(0)
+}
+
+func (m *MockUserDataStore) Put(ctx context.Context, key, value []byte) error {
+	args := m.Called(ctx, key, value)
+	return args.Error(0)
+}
+func (m *MockUserDataStore) Close() error {
+	args := m.Called(nil)
+	return args.Error(0)
+}

internal/utils/userStore.go

@@ -0,0 +1,81 @@
+package utils
+
+import (
+	"context"
+
+	"git.defalsify.org/vise.git/db"
+	"git.defalsify.org/vise.git/lang"
+)
+
+type DataStore interface {
+	SetPrefix(prefix uint8)
+	SetSession(sessionId string)
+	Get(ctx context.Context, key []byte) ([]byte, error)
+	ReadEntry(ctx context.Context, sessionId string, typ DataTyp) ([]byte, error)
+	WriteEntry(ctx context.Context, sessionId string, typ DataTyp, value []byte) error
+	Connect(ctx context.Context, connStr string) error
+	SetLanguage(*lang.Language)
+	Close() error
+	Prefix() uint8
+	Put(ctx context.Context, key []byte, val []byte) error
+	Safe() bool
+	SetLock(typ uint8, locked bool) error
+}
+
+type UserDataStore struct {
+	Store db.Db
+}
+
+func (store UserDataStore) SetPrefix(prefix uint8) {
+	store.Store.SetPrefix(prefix)
+}
+
+func (store UserDataStore) SetLanguage(lang *lang.Language) {
+	store.Store.SetLanguage(lang)
+}
+
+func (store UserDataStore) SetLock(typ uint8, locked bool) error {
+	return store.Store.SetLock(typ, locked)
+}
+func (store UserDataStore) Safe() bool {
+	return store.Store.Safe()
+}
+
+func (store UserDataStore) Put(ctx context.Context, key []byte, val []byte) error {
+	return store.Store.Put(ctx, key, val)
+}
+
+func (store UserDataStore) Connect(ctx context.Context, connectionStr string) error {
+	return store.Store.Connect(ctx, connectionStr)
+}
+
+func (store UserDataStore) Close() error {
+	return store.Store.Close()
+}
+
+func (store UserDataStore) Prefix() uint8 {
+	return store.Store.Prefix()
+}
+
+func (store UserDataStore) SetSession(sessionId string) {
+	store.Store.SetSession(sessionId)
+}
+
+func (store UserDataStore) Get(ctx context.Context, key []byte) ([]byte, error) {
+	return store.Store.Get(ctx, key)
+}
+
+// ReadEntry retrieves an entry from the store based on the provided parameters.
+func (store UserDataStore) ReadEntry(ctx context.Context, sessionId string, typ DataTyp) ([]byte, error) {
+	store.Store.SetPrefix(db.DATATYPE_USERDATA)
+	store.Store.SetSession(sessionId)
+	k := PackKey(typ, []byte(sessionId))
+	return store.Get(ctx, k)
+}
+
+func (store UserDataStore) WriteEntry(ctx context.Context, sessionId string, typ DataTyp, value []byte) error {
+	store.Store.SetPrefix(db.DATATYPE_USERDATA)
+	store.Store.SetSession(sessionId)
+	k := PackKey(typ, []byte(sessionId))
+	return store.Store.Put(ctx, k, value)
+}