fix tests

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>

management/server/sql_store_test.go

@@ -68,17 +68,23 @@ func TestSqlite_SaveAccount_Large(t *testing.T) {
 func runLargeTest(t *testing.T, store Store) {
 	t.Helper()
 
-	account := newAccountWithId(context.Background(), "account_id", "testuser", "")
-	groupALL, err := account.GetGroupAll()
-	if err != nil {
-		t.Fatal(err)
-	}
+	accountID := "account_id"
+
+	err := newAccountWithId(context.Background(), store, accountID, "testuser", "")
+	assert.NoError(t, err, "failed to create account")
+
+	groupAll, err := store.GetGroupByName(context.Background(), LockingStrengthShare, accountID, "All")
+	assert.NoError(t, err, "failed to get group All")
+
 	setupKey, _ := GenerateDefaultSetupKey()
-	account.SetupKeys[setupKey.Key] = setupKey
+	setupKey.AccountID = accountID
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	assert.NoError(t, err, "failed to save setup key")
+
 	const numPerAccount = 6000
 	for n := 0; n < numPerAccount; n++ {
 		netIP := randomIPv4()
-		peerID := fmt.Sprintf("%s-peer-%d", account.Id, n)
+		peerID := fmt.Sprintf("%s-peer-%d", accountID, n)
 
 		peer := &nbpeer.Peer{
 			ID:         peerID,
@@ -90,16 +96,21 @@ func runLargeTest(t *testing.T, store Store) {
 			Status:     &nbpeer.PeerStatus{Connected: false, LastSeen: time.Now()},
 			SSHEnabled: false,
 		}
-		account.Peers[peerID] = peer
-		group, _ := account.GetGroupAll()
-		group.Peers = append(group.Peers, peerID)
-		user := &User{
-			Id:        fmt.Sprintf("%s-user-%d", account.Id, n),
-			AccountID: account.Id,
-		}
-		account.Users[user.Id] = user
+		err = store.SavePeer(context.Background(), LockingStrengthUpdate, accountID, peer)
+		assert.NoError(t, err, "failed to save peer")
+
+		err = store.AddPeerToAllGroup(context.Background(), accountID, peerID)
+		assert.NoError(t, err, "failed to add peer to all group")
+
+		err = store.SaveUser(context.Background(), LockingStrengthUpdate, &User{
+			Id:        fmt.Sprintf("%s-user-%d", accountID, n),
+			AccountID: accountID,
+		})
+		assert.NoError(t, err, "failed to save user")
+
 		route := &route2.Route{
 			ID:          route2.ID(fmt.Sprintf("network-id-%d", n)),
+			AccountID:   accountID,
 			Description: "base route",
 			NetID:       route2.NetID(fmt.Sprintf("network-id-%d", n)),
 			Network:     netip.MustParsePrefix(netIP.String() + "/24"),
@@ -107,22 +118,24 @@ func runLargeTest(t *testing.T, store Store) {
 			Metric:      9999,
 			Masquerade:  false,
 			Enabled:     true,
-			Groups:      []string{groupALL.ID},
+			Groups:      []string{groupAll.ID},
 		}
-		account.Routes[route.ID] = route
+		err = store.SaveRoute(context.Background(), LockingStrengthUpdate, route)
+		assert.NoError(t, err, "failed to save route")
 
-		group = &nbgroup.Group{
+		group := &nbgroup.Group{
 			ID:        fmt.Sprintf("group-id-%d", n),
-			AccountID: account.Id,
+			AccountID: accountID,
 			Name:      fmt.Sprintf("group-id-%d", n),
 			Issued:    "api",
 			Peers:     nil,
 		}
-		account.Groups[group.ID] = group
+		err = store.SaveGroup(context.Background(), LockingStrengthUpdate, group)
+		assert.NoError(t, err, "failed to save group")
 
 		nameserver := &nbdns.NameServerGroup{
 			ID:                   fmt.Sprintf("nameserver-id-%d", n),
-			AccountID:            account.Id,
+			AccountID:            accountID,
 			Name:                 fmt.Sprintf("nameserver-id-%d", n),
 			Description:          "",
 			NameServers:          []nbdns.NameServer{{IP: netip.MustParseAddr(netIP.String()), NSType: nbdns.UDPNameServerType}},
@@ -132,20 +145,20 @@ func runLargeTest(t *testing.T, store Store) {
 			Enabled:              false,
 			SearchDomainsEnabled: false,
 		}
-		account.NameServerGroups[nameserver.ID] = nameserver
+		err = store.SaveNameServerGroup(context.Background(), LockingStrengthUpdate, nameserver)
+		assert.NoError(t, err, "failed to save nameserver group")
 
-		setupKey, _ := GenerateDefaultSetupKey()
-		account.SetupKeys[setupKey.Key] = setupKey
+		setupKey, _ = GenerateDefaultSetupKey()
+		setupKey.AccountID = accountID
+		err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+		assert.NoError(t, err, "failed to save setup key")
 	}
 
-	err = store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
-
 	if len(store.GetAllAccounts(context.Background())) != 1 {
 		t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
 	}
 
-	a, err := store.GetAccount(context.Background(), account.Id)
+	a, err := store.GetAccount(context.Background(), accountID)
 	if a == nil {
 		t.Errorf("expecting Account to be stored after SaveAccount(): %v", err)
 	}
@@ -213,41 +226,49 @@ func TestSqlite_SaveAccount(t *testing.T) {
 	t.Setenv("NETBIRD_STORE_ENGINE", string(SqliteStoreEngine))
 	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
 	t.Cleanup(cleanUp)
-	assert.NoError(t, err)
+	require.NoError(t, err)
+
+	accountID := "account_id"
+	err = newAccountWithId(context.Background(), store, accountID, "testuser", "")
+	require.NoError(t, err, "failed to create account")
 
-	account := newAccountWithId(context.Background(), "account_id", "testuser", "")
 	setupKey, _ := GenerateDefaultSetupKey()
-	account.SetupKeys[setupKey.Key] = setupKey
-	account.Peers["testpeer"] = &nbpeer.Peer{
+	setupKey.AccountID = accountID
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	require.NoError(t, err, "failed to save setup key")
+
+	err = store.SavePeer(context.Background(), LockingStrengthUpdate, accountID, &nbpeer.Peer{
 		Key:    "peerkey",
 		IP:     net.IP{127, 0, 0, 1},
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-	}
+	})
+	require.NoError(t, err, "failed to save peer")
 
-	err = store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
+	accountID2 := "account_id2"
+	err = newAccountWithId(context.Background(), store, accountID2, "testuser2", "")
+	require.NoError(t, err, "failed to create account")
 
-	account2 := newAccountWithId(context.Background(), "account_id2", "testuser2", "")
 	setupKey, _ = GenerateDefaultSetupKey()
-	account2.SetupKeys[setupKey.Key] = setupKey
-	account2.Peers["testpeer2"] = &nbpeer.Peer{
+	setupKey.AccountID = accountID2
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	require.NoError(t, err, "failed to save setup key")
+
+	err = store.SavePeer(context.Background(), LockingStrengthUpdate, accountID2, &nbpeer.Peer{
 		Key:    "peerkey2",
 		IP:     net.IP{127, 0, 0, 2},
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name 2",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-	}
-
-	err = store.SaveAccount(context.Background(), account2)
-	require.NoError(t, err)
+	})
+	require.NoError(t, err, "failed to save peer")
 
 	if len(store.GetAllAccounts(context.Background())) != 2 {
 		t.Errorf("expecting 2 Accounts to be stored after SaveAccount()")
 	}
 
-	a, err := store.GetAccount(context.Background(), account.Id)
+	a, err := store.GetAccount(context.Background(), accountID)
 	if a == nil {
 		t.Errorf("expecting Account to be stored after SaveAccount(): %v", err)
 	}
@@ -288,36 +309,56 @@ func TestSqlite_DeleteAccount(t *testing.T) {
 	t.Cleanup(cleanUp)
 	assert.NoError(t, err)
 
+	accountID := "account_id"
 	testUserID := "testuser"
+
 	user := NewAdminUser(testUserID)
 	user.PATs = map[string]*PersonalAccessToken{"testtoken": {
 		ID:   "testtoken",
 		Name: "test token",
 	}}
 
-	account := newAccountWithId(context.Background(), "account_id", testUserID, "")
+	err = newAccountWithId(context.Background(), store, accountID, testUserID, "")
+	require.NoError(t, err)
+
 	setupKey, _ := GenerateDefaultSetupKey()
-	account.SetupKeys[setupKey.Key] = setupKey
-	account.Peers["testpeer"] = &nbpeer.Peer{
+	setupKey.AccountID = accountID
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	require.NoError(t, err, "failed to save setup key")
+
+	err = store.SavePeer(context.Background(), LockingStrengthUpdate, accountID, &nbpeer.Peer{
 		Key:    "peerkey",
 		IP:     net.IP{127, 0, 0, 1},
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-	}
-	account.Users[testUserID] = user
+	})
+	require.NoError(t, err, "failed to save peer")
 
-	err = store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
+	err = store.SavePAT(context.Background(), LockingStrengthUpdate, &PersonalAccessToken{
+		ID:     "testtoken",
+		UserID: testUserID,
+		Name:   "test token",
+	})
+	require.NoError(t, err, "failed to save personal access token")
 
-	if len(store.GetAllAccounts(context.Background())) != 1 {
+	accountIDs, err := store.GetAllAccountIDs(context.Background(), LockingStrengthShare)
+	require.NoError(t, err, "failed to get all account ids")
+
+	if len(accountIDs) != 1 {
 		t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
 	}
 
+	account, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err, "failed to get account")
+
 	err = store.DeleteAccount(context.Background(), account)
 	require.NoError(t, err)
 
-	if len(store.GetAllAccounts(context.Background())) != 0 {
+	accountIDs, err = store.GetAllAccountIDs(context.Background(), LockingStrengthShare)
+	require.NoError(t, err, "failed to get all account ids after DeleteAccount()")
+
+	if len(accountIDs) != 0 {
 		t.Errorf("expecting 0 Accounts to be stored after DeleteAccount()")
 	}
 
@@ -714,19 +755,28 @@ func newSqliteStore(t *testing.T) *SqlStore {
 }
 
 func newAccount(store Store, id int) error {
-	str := fmt.Sprintf("%s-%d", uuid.New().String(), id)
-	account := newAccountWithId(context.Background(), str, str+"-testuser", "example.com")
+	accountID := fmt.Sprintf("%s-%d", uuid.New().String(), id)
+	userID := accountID + "-testuser"
+
+	err := newAccountWithId(context.Background(), store, accountID, userID, "example.com")
+	if err != nil {
+		return err
+	}
+
 	setupKey, _ := GenerateDefaultSetupKey()
-	account.SetupKeys[setupKey.Key] = setupKey
-	account.Peers["p"+str] = &nbpeer.Peer{
-		Key:    "peerkey" + str,
+	setupKey.AccountID = accountID
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	if err != nil {
+		return err
+	}
+
+	return store.SavePeer(context.Background(), LockingStrengthUpdate, accountID, &nbpeer.Peer{
+		Key:    accountID + "-peerkey",
 		IP:     net.IP{127, 0, 0, 1},
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-	}
-
-	return store.SaveAccount(context.Background(), account)
+	})
 }
 
 func TestPostgresql_NewStore(t *testing.T) {
@@ -754,39 +804,52 @@ func TestPostgresql_SaveAccount(t *testing.T) {
 	t.Cleanup(cleanUp)
 	assert.NoError(t, err)
 
-	account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+	accountID := "account_id"
+
+	err = newAccountWithId(context.Background(), store, accountID, "testuser", "")
+	require.NoError(t, err, "failed to create account")
+
 	setupKey, _ := GenerateDefaultSetupKey()
-	account.SetupKeys[setupKey.Key] = setupKey
-	account.Peers["testpeer"] = &nbpeer.Peer{
+	setupKey.AccountID = accountID
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	require.NoError(t, err, "failed to save setup key")
+
+	err = store.SavePeer(context.Background(), LockingStrengthUpdate, accountID, &nbpeer.Peer{
 		Key:    "peerkey",
 		IP:     net.IP{127, 0, 0, 1},
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-	}
+	})
+	require.NoError(t, err, "failed to save peer")
 
-	err = store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
+	accountID2 := "account_id2"
+
+	err = newAccountWithId(context.Background(), store, accountID2, "testuser2", "")
+	require.NoError(t, err, "failed to create account")
 
-	account2 := newAccountWithId(context.Background(), "account_id2", "testuser2", "")
 	setupKey, _ = GenerateDefaultSetupKey()
-	account2.SetupKeys[setupKey.Key] = setupKey
-	account2.Peers["testpeer2"] = &nbpeer.Peer{
+	setupKey.AccountID = accountID2
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	require.NoError(t, err, "failed to save setup key")
+
+	err = store.SavePeer(context.Background(), LockingStrengthUpdate, accountID2, &nbpeer.Peer{
 		Key:    "peerkey2",
 		IP:     net.IP{127, 0, 0, 2},
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name 2",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-	}
+	})
+	require.NoError(t, err, "failed to save peer")
 
-	err = store.SaveAccount(context.Background(), account2)
-	require.NoError(t, err)
+	accountIDs, err := store.GetAllAccountIDs(context.Background(), LockingStrengthUpdate)
+	require.NoError(t, err, "failed to get all account ids")
 
-	if len(store.GetAllAccounts(context.Background())) != 2 {
+	if len(accountIDs) != 2 {
 		t.Errorf("expecting 2 Accounts to be stored after SaveAccount()")
 	}
 
-	a, err := store.GetAccount(context.Background(), account.Id)
+	a, err := store.GetAccount(context.Background(), accountID)
 	if a == nil {
 		t.Errorf("expecting Account to be stored after SaveAccount(): %v", err)
 	}
@@ -827,32 +890,49 @@ func TestPostgresql_DeleteAccount(t *testing.T) {
 	t.Cleanup(cleanUp)
 	assert.NoError(t, err)
 
+	accountID := "account_id"
 	testUserID := "testuser"
+
 	user := NewAdminUser(testUserID)
 	user.PATs = map[string]*PersonalAccessToken{"testtoken": {
 		ID:   "testtoken",
 		Name: "test token",
 	}}
 
-	account := newAccountWithId(context.Background(), "account_id", testUserID, "")
+	err = newAccountWithId(context.Background(), store, accountID, testUserID, "")
+	require.NoError(t, err, "failed to create account")
+
 	setupKey, _ := GenerateDefaultSetupKey()
-	account.SetupKeys[setupKey.Key] = setupKey
-	account.Peers["testpeer"] = &nbpeer.Peer{
+	setupKey.AccountID = accountID
+	err = store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	require.NoError(t, err, "failed to save setup key")
+
+	err = store.SavePeer(context.Background(), LockingStrengthUpdate, accountID, &nbpeer.Peer{
 		Key:    "peerkey",
 		IP:     net.IP{127, 0, 0, 1},
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-	}
-	account.Users[testUserID] = user
+	})
+	require.NoError(t, err, "failed to save peer")
 
-	err = store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
+	err = store.SavePAT(context.Background(), LockingStrengthUpdate, &PersonalAccessToken{
+		ID:     "testtoken",
+		UserID: testUserID,
+		Name:   "test token",
+	})
+	require.NoError(t, err, "failed to save personal access token")
 
-	if len(store.GetAllAccounts(context.Background())) != 1 {
+	accountIDs, err := store.GetAllAccountIDs(context.Background(), LockingStrengthUpdate)
+	require.NoError(t, err, "failed to get all account ids")
+
+	if len(accountIDs) != 1 {
 		t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
 	}
 
+	account, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err, "failed to get account")
+
 	err = store.DeleteAccount(context.Background(), account)
 	require.NoError(t, err)
 
@@ -1218,7 +1298,7 @@ func TestSqlStore_UpdateAccountDomainAttributes(t *testing.T) {
 		domain := "example.com"
 		category := "public"
 		IsDomainPrimaryAccount := false
-		err = store.UpdateAccountDomainAttributes(context.Background(), accountID, domain, category, IsDomainPrimaryAccount)
+		err = store.UpdateAccountDomainAttributes(context.Background(), LockingStrengthUpdate, accountID, domain, category, &IsDomainPrimaryAccount)
 		require.NoError(t, err)
 		account, err := store.GetAccount(context.Background(), accountID)
 		require.NoError(t, err)
@@ -1232,7 +1312,7 @@ func TestSqlStore_UpdateAccountDomainAttributes(t *testing.T) {
 		domain := "test.com"
 		category := "private"
 		IsDomainPrimaryAccount := true
-		err = store.UpdateAccountDomainAttributes(context.Background(), accountID, domain, category, IsDomainPrimaryAccount)
+		err = store.UpdateAccountDomainAttributes(context.Background(), LockingStrengthUpdate, accountID, domain, category, &IsDomainPrimaryAccount)
 		require.NoError(t, err)
 		account, err := store.GetAccount(context.Background(), accountID)
 		require.NoError(t, err)
@@ -1246,7 +1326,9 @@ func TestSqlStore_UpdateAccountDomainAttributes(t *testing.T) {
 		domain := "test.com"
 		category := "private"
 		IsDomainPrimaryAccount := true
-		err = store.UpdateAccountDomainAttributes(context.Background(), "non-existing-account-id", domain, category, IsDomainPrimaryAccount)
+		err = store.UpdateAccountDomainAttributes(context.Background(), LockingStrengthUpdate, "non-existing-account-id",
+			domain, category, &IsDomainPrimaryAccount,
+		)
 		require.Error(t, err)
 	})
 
@@ -1274,7 +1356,7 @@ func Test_DeleteSetupKeySuccessfully(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	setupKeyID := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
 
-	err = store.DeleteSetupKey(context.Background(), accountID, setupKeyID)
+	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, setupKeyID)
 	require.NoError(t, err)
 
 	_, err = store.GetSetupKeyByID(context.Background(), LockingStrengthShare, setupKeyID, accountID)
@@ -1290,6 +1372,6 @@ func Test_DeleteSetupKeyFailsForNonExistingKey(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	nonExistingKeyID := "non-existing-key-id"
 
-	err = store.DeleteSetupKey(context.Background(), accountID, nonExistingKeyID)
+	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, nonExistingKeyID)
 	require.Error(t, err)
 }