Fix tests

Signed-off-by: bcmmbaga <[email protected]>
netbirdio · Nov 14, 2024 · f6f7260 · f6f7260
1 parent c557c98
commit f6f7260
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 23 deletions.
diff --git a/management/server/account_test.go b/management/server/account_test.go
@@ -1472,7 +1472,6 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 		return
 	}
 
-	userID := "account_creator"
 	account, err := createAccount(manager, "test_account", userID, "netbird.cloud")
 	if err != nil {
 		t.Fatal(err)
@@ -1501,7 +1500,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 		return
 	}
 
-	err = manager.DeletePeer(context.Background(), account.Id, peerKey, userID)
+	err = manager.DeletePeer(context.Background(), account.Id, peer.ID, userID)
 	if err != nil {
 		return
 	}
@@ -1523,7 +1522,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 	assert.Equal(t, peer.Name, ev.Meta["name"])
 	assert.Equal(t, peer.FQDN(account.Domain), ev.Meta["fqdn"])
 	assert.Equal(t, userID, ev.InitiatorID)
-	assert.Equal(t, peer.IP.String(), ev.TargetID)
+	assert.Equal(t, peer.ID, ev.TargetID)
 	assert.Equal(t, peer.IP.String(), fmt.Sprint(ev.Meta["ip"]))
 }
 
@@ -1853,13 +1852,10 @@ func TestDefaultAccountManager_UpdatePeer_PeerLoginExpiration(t *testing.T) {
 	accountID, err := manager.GetAccountIDByUserID(context.Background(), userID, "")
 	require.NoError(t, err, "unable to get the account")
 
-	account, err := manager.Store.GetAccount(context.Background(), accountID)
-	require.NoError(t, err, "unable to get the account")
-
-	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, account)
+	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, accountID)
 	require.NoError(t, err, "unable to mark peer connected")
 
-	account, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &Settings{
+	account, err := manager.UpdateAccountSettings(context.Background(), accountID, userID, &Settings{
 		PeerLoginExpiration:        time.Hour,
 		PeerLoginExpirationEnabled: true,
 	})
@@ -1927,11 +1923,8 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
 	accountID, err = manager.GetAccountIDByUserID(context.Background(), userID, "")
 	require.NoError(t, err, "unable to get the account")
 
-	account, err := manager.Store.GetAccount(context.Background(), accountID)
-	require.NoError(t, err, "unable to get the account")
-
 	// when we mark peer as connected, the peer login expiration routine should trigger
-	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, account)
+	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, accountID)
 	require.NoError(t, err, "unable to mark peer connected")
 
 	failed := waitTimeout(wg, time.Second)
@@ -1962,7 +1955,7 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test
 	account, err := manager.Store.GetAccount(context.Background(), accountID)
 	require.NoError(t, err, "unable to get the account")
 
-	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, account)
+	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, accountID)
 	require.NoError(t, err, "unable to mark peer connected")
 
 	wg := &sync.WaitGroup{}

diff --git a/management/server/sql_store_test.go b/management/server/sql_store_test.go
@@ -400,7 +400,7 @@ func TestSqlite_SavePeer(t *testing.T) {
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
 	}
 	ctx := context.Background()
-	err = store.SavePeer(ctx, account.Id, peer)
+	err = store.SavePeer(ctx, LockingStrengthUpdate, account.Id, peer)
 	assert.Error(t, err)
 	parsedErr, ok := status.FromError(err)
 	require.True(t, ok)
@@ -416,7 +416,7 @@ func TestSqlite_SavePeer(t *testing.T) {
 	updatedPeer.Status.Connected = false
 	updatedPeer.Meta.Hostname = "updatedpeer"
 
-	err = store.SavePeer(ctx, account.Id, updatedPeer)
+	err = store.SavePeer(ctx, LockingStrengthUpdate, account.Id, updatedPeer)
 	require.NoError(t, err)
 
 	account, err = store.GetAccount(context.Background(), account.Id)
@@ -442,7 +442,7 @@ func TestSqlite_SavePeerStatus(t *testing.T) {
 
 	// save status of non-existing peer
 	newStatus := nbpeer.PeerStatus{Connected: false, LastSeen: time.Now().UTC()}
-	err = store.SavePeerStatus(account.Id, "non-existing-peer", newStatus)
+	err = store.SavePeerStatus(context.Background(), LockingStrengthUpdate, account.Id, "non-existing-peer", newStatus)
 	assert.Error(t, err)
 	parsedErr, ok := status.FromError(err)
 	require.True(t, ok)
@@ -461,7 +461,7 @@ func TestSqlite_SavePeerStatus(t *testing.T) {
 	err = store.SaveAccount(context.Background(), account)
 	require.NoError(t, err)
 
-	err = store.SavePeerStatus(account.Id, "testpeer", newStatus)
+	err = store.SavePeerStatus(context.Background(), LockingStrengthUpdate, account.Id, "testpeer", newStatus)
 	require.NoError(t, err)
 
 	account, err = store.GetAccount(context.Background(), account.Id)
@@ -472,7 +472,7 @@ func TestSqlite_SavePeerStatus(t *testing.T) {
 
 	newStatus.Connected = true
 
-	err = store.SavePeerStatus(account.Id, "testpeer", newStatus)
+	err = store.SavePeerStatus(context.Background(), LockingStrengthUpdate, account.Id, "testpeer", newStatus)
 	require.NoError(t, err)
 
 	account, err = store.GetAccount(context.Background(), account.Id)
@@ -507,7 +507,7 @@ func TestSqlite_SavePeerLocation(t *testing.T) {
 		Meta: nbpeer.PeerSystemMeta{},
 	}
 	// error is expected as peer is not in store yet
-	err = store.SavePeerLocation(account.Id, peer)
+	err = store.SavePeerLocation(context.Background(), LockingStrengthUpdate, account.Id, peer)
 	assert.Error(t, err)
 
 	account.Peers[peer.ID] = peer
@@ -519,7 +519,7 @@ func TestSqlite_SavePeerLocation(t *testing.T) {
 	peer.Location.CityName = "Berlin"
 	peer.Location.GeoNameID = 2950159
 
-	err = store.SavePeerLocation(account.Id, account.Peers[peer.ID])
+	err = store.SavePeerLocation(context.Background(), LockingStrengthUpdate, account.Id, account.Peers[peer.ID])
 	assert.NoError(t, err)
 
 	account, err = store.GetAccount(context.Background(), account.Id)
@@ -529,7 +529,7 @@ func TestSqlite_SavePeerLocation(t *testing.T) {
 	assert.Equal(t, peer.Location, actual)
 
 	peer.ID = "non-existing-peer"
-	err = store.SavePeerLocation(account.Id, peer)
+	err = store.SavePeerLocation(context.Background(), LockingStrengthUpdate, account.Id, peer)
 	assert.Error(t, err)
 	parsedErr, ok := status.FromError(err)
 	require.True(t, ok)
@@ -908,7 +908,7 @@ func TestPostgresql_SavePeerStatus(t *testing.T) {
 
 	// save status of non-existing peer
 	newStatus := nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()}
-	err = store.SavePeerStatus(account.Id, "non-existing-peer", newStatus)
+	err = store.SavePeerStatus(context.Background(), LockingStrengthUpdate, account.Id, "non-existing-peer", newStatus)
 	assert.Error(t, err)
 
 	// save new status of existing peer
@@ -924,7 +924,7 @@ func TestPostgresql_SavePeerStatus(t *testing.T) {
 	err = store.SaveAccount(context.Background(), account)
 	require.NoError(t, err)
 
-	err = store.SavePeerStatus(account.Id, "testpeer", newStatus)
+	err = store.SavePeerStatus(context.Background(), LockingStrengthUpdate, account.Id, "testpeer", newStatus)
 	require.NoError(t, err)
 
 	account, err = store.GetAccount(context.Background(), account.Id)