diff --git a/tools/walletextension/storage/storage_test.go b/tools/walletextension/storage/storage_test.go
index 1913555007..9fde210cc6 100644
--- a/tools/walletextension/storage/storage_test.go
+++ b/tools/walletextension/storage/storage_test.go
@@ -6,6 +6,8 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/ten-protocol/go-ten/integration/common/testlog"
+
 	"github.com/ten-protocol/go-ten/go/common/viewingkey"
 
 	"github.com/stretchr/testify/require"
@@ -26,7 +28,7 @@ func TestGatewayStorage(t *testing.T) {
 
 	for name, test := range tests {
 		t.Run(name, func(t *testing.T) {
-			storage, err := New("sqlite", "", "", randomKey)
+			storage, err := New("sqlite", "", "", randomKey, testlog.Logger())
 			// storage, err := New("cosmosDB", "", "", randomKey)
 			require.NoError(t, err)
 
@@ -61,8 +63,8 @@ func testAddAndGetUser(storage UserStorage, t *testing.T) {
 	}
 
 	// Check if retrieved private key matches the original
-	if !bytes.Equal(user.PrivateKey, privateKey) {
-		t.Errorf("privateKey mismatch: got %v, want %v", user.PrivateKey, privateKey)
+	if !bytes.Equal(user.UserKey, privateKey) {
+		t.Errorf("privateKey mismatch: got %v, want %v", user.UserKey, privateKey)
 	}
 }
 
@@ -118,10 +120,10 @@ func testAddAccounts(storage UserStorage, t *testing.T) {
 
 	// Iterate through retrieved accounts and check if they match the added accounts
 	for _, account := range user.Accounts {
-		if bytes.Equal(account.AccountAddress, accountAddress1) && bytes.Equal(account.Signature, signature1) {
+		if bytes.Equal(account.Address.Bytes(), accountAddress1) && bytes.Equal(account.Signature, signature1) {
 			foundAccount1 = true
 		}
-		if bytes.Equal(account.AccountAddress, accountAddress2) && bytes.Equal(account.Signature, signature2) {
+		if bytes.Equal(account.Address.Bytes(), accountAddress2) && bytes.Equal(account.Signature, signature2) {
 			foundAccount2 = true
 		}
 	}
@@ -183,12 +185,12 @@ func testGetUser(storage UserStorage, t *testing.T) {
 	}
 
 	// Check if retrieved user matches the added user
-	if !bytes.Equal(user.UserId, userID) {
-		t.Errorf("Retrieved user ID does not match. Expected %x, got %x", userID, user.UserId)
+	if !bytes.Equal(user.UserID, userID) {
+		t.Errorf("Retrieved user ID does not match. Expected %x, got %x", userID, user.UserID)
 	}
 
-	if !bytes.Equal(user.PrivateKey, privateKey) {
-		t.Errorf("Retrieved private key does not match. Expected %x, got %x", privateKey, user.PrivateKey)
+	if !bytes.Equal(user.UserKey, privateKey) {
+		t.Errorf("Retrieved private key does not match. Expected %x, got %x", privateKey, user.UserKey)
 	}
 
 	// Try to get a non-existent user