feat(authz): update integration tests

pkg/modules/serviceaccount/implserviceaccount/module.go

@@ -376,7 +376,7 @@ func (module *module) getOrGetSetIdentity(ctx context.Context, serviceAccountID
 }
 
 func (module *module) setRole(ctx context.Context, orgID valuer.UUID, id valuer.UUID, role *authtypes.Role) error {
-	serviceAccount, err := module.Get(ctx, orgID, id)
+	serviceAccount, err := module.GetWithRoles(ctx, orgID, id)
 	if err != nil {
 		return err
 	}
@@ -386,12 +386,24 @@ func (module *module) setRole(ctx context.Context, orgID valuer.UUID, id valuer.
 		return err
 	}
 
-	err = module.authz.Grant(ctx, orgID, []string{role.Name}, authtypes.MustNewSubject(authtypes.TypeableServiceAccount, id.String(), orgID, nil))
+	err = module.authz.ModifyGrant(ctx, orgID, serviceAccount.RoleNames(), []string{role.Name}, authtypes.MustNewSubject(authtypes.TypeableServiceAccount, id.String(), orgID, nil))
 	if err != nil {
 		return err
 	}
 
-	err = module.store.CreateServiceAccountRole(ctx, serviceAccountRole)
+	err = module.store.RunInTx(ctx, func(ctx context.Context) error {
+		err = module.store.DeleteServiceAccountRoles(ctx, serviceAccount.ID)
+		if err != nil {
+			return err
+		}
+
+		err = module.store.CreateServiceAccountRole(ctx, serviceAccountRole)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
 	if err != nil {
 		return err
 	}

pkg/modules/serviceaccount/implserviceaccount/store.go

@@ -170,6 +170,21 @@ func (store *store) CreateServiceAccountRole(ctx context.Context, serviceAccount
 	return nil
 }
 
+func (store *store) DeleteServiceAccountRoles(ctx context.Context, serviceAccountID valuer.UUID) error {
+	_, err := store.
+		sqlstore.
+		BunDBCtx(ctx).
+		NewDelete().
+		Model(new(serviceaccounttypes.ServiceAccountRole)).
+		Where("service_account_id = ?", serviceAccountID).
+		Exec(ctx)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (store *store) DeleteServiceAccountRole(ctx context.Context, serviceAccountID valuer.UUID, roleID valuer.UUID) error {
 	_, err := store.
 		sqlstore.

pkg/modules/user/impluser/setter.go

@@ -866,16 +866,17 @@ func (module *setter) AddUserRole(ctx context.Context, orgID, userID valuer.UUID
 	if err != nil {
 		return err
 	}
-	for _, userRole := range existingUserRoles {
-		if userRole.Role != nil && userRole.Role.Name == roleName {
-			return nil // role already assigned no-op
-		}
+
+	existingRoles := make([]string, len(existingUserRoles))
+	for idx, role := range existingUserRoles {
+		existingRoles[idx] = role.Role.Name
 	}
 
 	// grant via authz (idempotent)
-	if err := module.authz.Grant(
+	if err := module.authz.ModifyGrant(
 		ctx,
 		orgID,
+		existingRoles,
 		[]string{roleName},
 		authtypes.MustNewSubject(authtypes.TypeableUser, existingUser.ID.StringValue(), existingUser.OrgID, nil),
 	); err != nil {
@@ -884,7 +885,20 @@ func (module *setter) AddUserRole(ctx context.Context, orgID, userID valuer.UUID
 
 	// create user_role entry
 	userRoles := authtypes.NewUserRoles(userID, foundRoles)
-	if err := module.userRoleStore.CreateUserRoles(ctx, userRoles); err != nil {
+	err = module.store.RunInTx(ctx, func(ctx context.Context) error {
+		err = module.userRoleStore.DeleteUserRoles(ctx, existingUser.ID)
+		if err != nil {
+			return err
+		}
+
+		err := module.userRoleStore.CreateUserRoles(ctx, userRoles)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+	if err != nil {
 		return err
 	}
 

pkg/types/serviceaccounttypes/service_acccount.go

@@ -243,6 +243,7 @@ type Store interface {
 
 	// Service Account Role
 	CreateServiceAccountRole(context.Context, *ServiceAccountRole) error
+	DeleteServiceAccountRoles(context.Context, valuer.UUID) error
 	DeleteServiceAccountRole(context.Context, valuer.UUID, valuer.UUID) error
 
 	// Service Account Factor API Key

tests/integration/src/passwordauthn/04_role.py

@@ -133,11 +133,11 @@ def test_get_user_roles(
         assert "type" in role
 
 
-def test_assign_additional_role(
+def test_assign_role_replaces_previous(
     signoz: types.SigNoz,
     get_token: Callable[[str, str], str],
 ):
-    """Verify POST /api/v2/users/{id}/roles assigns an additional role."""
+    """Verify POST /api/v2/users/{id}/roles replaces existing role."""
     admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
     response = requests.get(
         signoz.self.host_configs["8080"].get("/api/v2/users/me"),
@@ -166,8 +166,8 @@ def test_assign_additional_role(
     assert response.status_code == HTTPStatus.OK
     roles = response.json()["data"]
     names = {r["name"] for r in roles}
-    assert "signoz-admin" in names
     assert "signoz-editor" in names
+    assert "signoz-admin" not in names
 
 
 def test_get_users_by_role(
@@ -253,9 +253,7 @@ def test_remove_role(
     )
     assert response.status_code == HTTPStatus.OK
     roles_after = response.json()["data"]
-    names = {r["name"] for r in roles_after}
-    assert "signoz-editor" not in names
-    assert "signoz-admin" in names
+    assert len(roles_after) == 0
 
 
 def test_user_with_roles_reflects_change(
@@ -282,8 +280,7 @@ def test_user_with_roles_reflects_change(
     assert response.status_code == HTTPStatus.OK
     data = response.json()["data"]
     role_names = {ur["role"]["name"] for ur in data["userRoles"]}
-    assert "signoz-admin" in role_names
-    assert "signoz-editor" not in role_names
+    assert len(role_names) == 0
 
 
 def test_admin_cannot_assign_role_to_self(

tests/integration/src/serviceaccount/04_roles.py

@@ -48,7 +48,7 @@ def test_assign_role_to_service_account(
     create_user_admin: types.Operation,  # pylint: disable=unused-argument
     get_token: Callable[[str, str], str],
 ):
-    """POST /{id}/roles assigns a new role, verify via GET."""
+    """POST /{id}/roles replaces existing role, verify via GET."""
     token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
 
     # create service account with viewer role
@@ -56,7 +56,7 @@ def test_assign_role_to_service_account(
         signoz, token, "sa-assign-role", role="signoz-viewer"
     )
 
-    # assign editor role additionally
+    # assign editor role (replaces viewer)
     editor_role_id = find_role_by_name(signoz, token, "signoz-editor")
     assign_resp = requests.post(
         signoz.self.host_configs["8080"].get(
@@ -68,7 +68,7 @@ def test_assign_role_to_service_account(
     )
     assert assign_resp.status_code == HTTPStatus.NO_CONTENT, assign_resp.text
 
-    # verify both roles are present
+    # verify only editor role is present (viewer was replaced)
     roles_resp = requests.get(
         signoz.self.host_configs["8080"].get(
             f"{SERVICE_ACCOUNT_BASE}/{service_account_id}/roles"
@@ -78,8 +78,9 @@ def test_assign_role_to_service_account(
     )
     assert roles_resp.status_code == HTTPStatus.OK, roles_resp.text
     role_names = [r["name"] for r in roles_resp.json()["data"]]
-    assert "signoz-viewer" in role_names
+    assert len(role_names) == 1
     assert "signoz-editor" in role_names
+    assert "signoz-viewer" not in role_names
 
 
 def test_assign_role_idempotent(
@@ -87,7 +88,7 @@ def test_assign_role_idempotent(
     create_user_admin: types.Operation,  # pylint: disable=unused-argument
     get_token: Callable[[str, str], str],
 ):
-    """POST same role twice succeeds (store uses ON CONFLICT DO NOTHING)."""
+    """POST same role twice succeeds (replace with same role is idempotent)."""
     token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
     service_account_id = create_service_account(
         signoz, token, "sa-role-idempotent", role="signoz-viewer"
@@ -118,6 +119,66 @@ def test_assign_role_idempotent(
     assert role_names.count("signoz-viewer") == 1
 
 
+def test_assign_role_replaces_access(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+):
+    """After role replacement, SA loses old permissions and gains new ones."""
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    # create SA with viewer role and an API key
+    service_account_id, api_key = create_service_account_with_key(
+        signoz, token, "sa-role-replace-access", role="signoz-viewer"
+    )
+
+    # viewer should get 403 on admin-only endpoint
+    resp = requests.get(
+        signoz.self.host_configs["8080"].get(SERVICE_ACCOUNT_BASE),
+        headers={"SIGNOZ-API-KEY": api_key},
+        timeout=5,
+    )
+    assert (
+        resp.status_code == HTTPStatus.FORBIDDEN
+    ), f"Expected 403 for viewer on admin endpoint, got {resp.status_code}: {resp.text}"
+
+    # assign admin role (replaces viewer)
+    admin_role_id = find_role_by_name(signoz, token, "signoz-admin")
+    assign_resp = requests.post(
+        signoz.self.host_configs["8080"].get(
+            f"{SERVICE_ACCOUNT_BASE}/{service_account_id}/roles"
+        ),
+        json={"id": admin_role_id},
+        headers={"Authorization": f"Bearer {token}"},
+        timeout=5,
+    )
+    assert assign_resp.status_code == HTTPStatus.NO_CONTENT, assign_resp.text
+
+    # SA should now have admin access
+    resp = requests.get(
+        signoz.self.host_configs["8080"].get(SERVICE_ACCOUNT_BASE),
+        headers={"SIGNOZ-API-KEY": api_key},
+        timeout=5,
+    )
+    assert (
+        resp.status_code == HTTPStatus.OK
+    ), f"Expected 200 for admin on admin endpoint, got {resp.status_code}: {resp.text}"
+
+    # verify only admin role is present
+    roles_resp = requests.get(
+        signoz.self.host_configs["8080"].get(
+            f"{SERVICE_ACCOUNT_BASE}/{service_account_id}/roles"
+        ),
+        headers={"Authorization": f"Bearer {token}"},
+        timeout=5,
+    )
+    assert roles_resp.status_code == HTTPStatus.OK, roles_resp.text
+    role_names = [r["name"] for r in roles_resp.json()["data"]]
+    assert len(role_names) == 1
+    assert "signoz-admin" in role_names
+    assert "signoz-viewer" not in role_names
+
+
 def test_remove_role_from_service_account(
     signoz: types.SigNoz,
     create_user_admin: types.Operation,  # pylint: disable=unused-argument