diff --git a/management/server/account_test.go b/management/server/account_test.go
index a13b89f33..650e8de69 100644
--- a/management/server/account_test.go
+++ b/management/server/account_test.go
@@ -758,6 +758,7 @@ func TestAccountManager_GetAccountFromPAT(t *testing.T) {
 		PATs: map[string]*PersonalAccessToken{
 			"tokenId": {
 				ID:          "tokenId",
+				UserID:      "someUser",
 				HashedToken: encodedHashedToken,
 			},
 		},
@@ -771,14 +772,14 @@ func TestAccountManager_GetAccountFromPAT(t *testing.T) {
 		Store: store,
 	}
 
-	account, user, pat, err := am.GetAccountFromPAT(context.Background(), token)
+	user, pat, _, _, err := am.GetAccountInfoFromPAT(context.Background(), token)
 	if err != nil {
 		t.Fatalf("Error when getting Account from PAT: %s", err)
 	}
 
-	assert.Equal(t, "account_id", account.Id)
+	assert.Equal(t, "account_id", user.AccountID)
 	assert.Equal(t, "someUser", user.Id)
-	assert.Equal(t, account.Users["someUser"].PATs["tokenId"], pat)
+	assert.Equal(t, account.Users["someUser"].PATs["tokenId"].ID, pat.ID)
 }
 
 func TestDefaultAccountManager_MarkPATUsed(t *testing.T) {
diff --git a/management/server/peer_test.go b/management/server/peer_test.go
index e410fa892..0e30a3762 100644
--- a/management/server/peer_test.go
+++ b/management/server/peer_test.go
@@ -13,6 +13,7 @@ import (
 	"testing"
 	"time"
 
+	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
@@ -22,7 +23,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/domain"
 	"github.com/netbirdio/netbird/management/proto"
-	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
@@ -1446,7 +1446,8 @@ func TestPeerAccountPeersUpdate(t *testing.T) {
 	// Adding peer to group linked with policy should update account peers and send peer update
 	t.Run("adding peer to group linked with policy", func(t *testing.T) {
 		_, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
-			Enabled: true,
+			AccountID: account.Id,
+			Enabled:   true,
 			Rules: []*PolicyRule{
 				{
 					Enabled:       true,
diff --git a/management/server/sql_store_test.go b/management/server/sql_store_test.go
index 39c42bfa6..211767bff 100644
--- a/management/server/sql_store_test.go
+++ b/management/server/sql_store_test.go
@@ -589,11 +589,11 @@ func TestSqlite_GetUserByTokenID(t *testing.T) {
 
 	id := "9dj38s35-63fb-11ec-90d6-0242ac120003"
 
-	user, err := store.GetUserByTokenID(context.Background(), id)
+	user, err := store.GetUserByPATID(context.Background(), LockingStrengthShare, id)
 	require.NoError(t, err)
-	require.Equal(t, id, user.PATs[id].ID)
+	require.NotNil(t, id, user)
 
-	_, err = store.GetUserByTokenID(context.Background(), "non-existing-id")
+	_, err = store.GetUserByPATID(context.Background(), LockingStrengthShare, "non-existing-id")
 	require.Error(t, err)
 	parsedErr, ok := status.FromError(err)
 	require.True(t, ok)
@@ -978,7 +978,7 @@ func TestPostgresql_GetUserByTokenID(t *testing.T) {
 
 	id := "9dj38s35-63fb-11ec-90d6-0242ac120003"
 
-	user, err := store.GetUserByTokenID(context.Background(), id)
+	user, err := store.GetUserByPATID(context.Background(), LockingStrengthShare, id)
 	require.NoError(t, err)
 	require.Equal(t, id, user.PATs[id].ID)
 }
diff --git a/management/server/user_test.go b/management/server/user_test.go
index 498017afa..c599f6331 100644
--- a/management/server/user_test.go
+++ b/management/server/user_test.go
@@ -73,7 +73,7 @@ func TestUser_CreatePAT_ForSameUser(t *testing.T) {
 
 	assert.Equal(t, pat.ID, tokenID)
 
-	user, err := am.Store.GetUserByTokenID(context.Background(), tokenID)
+	user, err := am.Store.GetUserByPATID(context.Background(), LockingStrengthShare, tokenID)
 	if err != nil {
 		t.Fatalf("Error when getting user by token ID: %s", err)
 	}