Improve private domain's behavior tests and logic (#256)

Improved the behavior tests for private domains
and its logic as well because on existing accounts
there was no primary status update

management/server/account.go

@@ -248,11 +248,16 @@ func (am *DefaultAccountManager) updateAccountDomainAttributes(account *Account,
 func (am *DefaultAccountManager) handleExistingUserAccount(existingAcc *Account, domainAcc *Account, claims jwtclaims.AuthorizationClaims) error {
 	var err error
 
-	if domainAcc == nil || existingAcc.Id != domainAcc.Id {
+	if domainAcc != nil && existingAcc.Id != domainAcc.Id {
 		err = am.updateAccountDomainAttributes(existingAcc, claims, false)
 		if err != nil {
 			return err
 		}
+	} else {
+		err = am.updateAccountDomainAttributes(existingAcc, claims, true)
+		if err != nil {
+			return err
+		}
 	}
 
 	// we should register the account ID to this user's metadata in our IDP manager
@@ -269,24 +274,21 @@ func (am *DefaultAccountManager) handleExistingUserAccount(existingAcc *Account,
 func (am *DefaultAccountManager) handleNewUserAccount(domainAcc *Account, claims jwtclaims.AuthorizationClaims) (*Account, error) {
 	var (
 		account *Account
-		primaryAccount bool
+		err     error
 	)
 	lowerDomain := strings.ToLower(claims.Domain)
 	// if domain already has a primary account, add regular user
 	if domainAcc != nil {
 		account = domainAcc
 		account.Users[claims.UserId] = NewRegularUser(claims.UserId)
-		primaryAccount = false
 	} else {
 		account = NewAccount(claims.UserId, lowerDomain)
 		account.Users[claims.UserId] = NewAdminUser(claims.UserId)
-		primaryAccount = true
+		err = am.updateAccountDomainAttributes(account, claims, true)
-	}
-
-	err := am.updateAccountDomainAttributes(account, claims, primaryAccount)
 		if err != nil {
 			return nil, err
 		}
+	}
 
 	err = am.updateIDPMetadata(claims.UserId, account.Id)
 	if err != nil {

management/server/account_test.go

@@ -48,6 +48,7 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 		expectedMSG                 string
 		expectedUserRole            UserRole
 		expectedDomainCategory      string
+		expectedPrimaryDomainStatus bool
 	}
 
 	var (
@@ -73,6 +74,7 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 		expectedMSG:                 "account IDs shouldn't match",
 		expectedUserRole:            UserRoleAdmin,
 		expectedDomainCategory:      "",
+		expectedPrimaryDomainStatus: false,
 	}
 
 	initUnknown := defaultInitAccount
@@ -91,6 +93,7 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 		expectedMSG:                 "account IDs shouldn't match",
 		expectedUserRole:            UserRoleAdmin,
 		expectedDomainCategory:      "",
+		expectedPrimaryDomainStatus: false,
 	}
 
 	testCase3 := test{
@@ -105,6 +108,7 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 		expectedMSG:                 "account IDs shouldn't match",
 		expectedUserRole:            UserRoleAdmin,
 		expectedDomainCategory:      PrivateCategory,
+		expectedPrimaryDomainStatus: true,
 	}
 
 	privateInitAccount := defaultInitAccount
@@ -124,6 +128,7 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 		expectedMSG:                 "account IDs should match",
 		expectedUserRole:            UserRoleUser,
 		expectedDomainCategory:      PrivateCategory,
+		expectedPrimaryDomainStatus: true,
 	}
 
 	testCase5 := test{
@@ -138,6 +143,7 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 		expectedMSG:                 "account IDs should match",
 		expectedUserRole:            UserRoleAdmin,
 		expectedDomainCategory:      PrivateCategory,
+		expectedPrimaryDomainStatus: true,
 	}
 
 	testCase6 := test{
@@ -153,6 +159,7 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 		expectedMSG:                 "account IDs should match",
 		expectedUserRole:            UserRoleAdmin,
 		expectedDomainCategory:      PrivateCategory,
+		expectedPrimaryDomainStatus: true,
 	}
 	for _, testCase := range []test{testCase1, testCase2, testCase3, testCase4, testCase5, testCase6} {
 		t.Run(testCase.name, func(t *testing.T) {
@@ -177,8 +184,9 @@ func TestDefaultAccountManager_GetAccountWithAuthorizationClaims(t *testing.T) {
 
 			testCase.testingFunc(t, initAccount.Id, account.Id, testCase.expectedMSG)
 
-			require.EqualValues(t, testCase.expectedUserRole, account.Users[testCase.inputClaims.UserId].Role, "user role should match")
+			require.EqualValues(t, testCase.expectedUserRole, account.Users[testCase.inputClaims.UserId].Role, "expected user role should match")
-			require.EqualValues(t, testCase.expectedDomainCategory, account.DomainCategory, "account domain category should match")
+			require.EqualValues(t, testCase.expectedDomainCategory, account.DomainCategory, "expected account domain category should match")
+			require.EqualValues(t, testCase.expectedPrimaryDomainStatus, account.IsDomainPrimaryAccount, "expected account primary status should match")
 		})
 	}
 }