From 28562be0174f22f5d658c2cfd80c5b178f567091 Mon Sep 17 00:00:00 2001 From: Michelle Tran Date: Thu, 1 Aug 2024 14:55:30 -0400 Subject: [PATCH] Sync account when creating in Okta This just syncs a bunch of other fields for accounts created during the Okta settings creation process. --- .../owner/interactors/save_okta_config.py | 41 ++++++++++++++++--- .../tests/test_save_okta_config.py | 33 +++++++++++++-- 2 files changed, 65 insertions(+), 9 deletions(-) diff --git a/codecov_auth/commands/owner/interactors/save_okta_config.py b/codecov_auth/commands/owner/interactors/save_okta_config.py index f6b66d5cbf..14744568c8 100644 --- a/codecov_auth/commands/owner/interactors/save_okta_config.py +++ b/codecov_auth/commands/owner/interactors/save_okta_config.py @@ -1,6 +1,9 @@ from dataclasses import dataclass from typing import Optional +from django.contrib.auth.models import User +from shared.django_apps.codecov_auth.models import AccountsUsers + from codecov.commands.base import BaseInteractor from codecov.commands.exceptions import Unauthenticated, Unauthorized, ValidationError from codecov.db import sync_to_async @@ -42,14 +45,40 @@ def execute(self, input: dict): ).first() self.validate(owner=owner) - account_id = owner.account_id - if not account_id: - account = Account.objects.create() - account_id = account.id - owner.account_id = account_id + account = owner.account + if not account: + account = Account.objects.create( + name=owner.username, + plan=owner.plan, + plan_seat_count=owner.plan_user_count, + free_seat_count=owner.free, + plan_auto_activate=owner.plan_auto_activate, + ) + owner.account = account owner.save() - okta_config, created = OktaSettings.objects.get_or_create(account_id=account_id) + # Update the activated users to be added to the account + plan_activated_user_owners: list[int] = owner.plan_activated_users + activated_connections: list[AccountsUsers] = [] + for activated_user_owner in plan_activated_user_owners: + user_owner: Owner = Owner.objects.get(pk=activated_user_owner) + user = user_owner.user + if user is None: + user = User(name=user_owner.name, email=user_owner.email) + user_owner.userid = user.id + user_owner.save() + + activated_connections.append(AccountsUsers(account=account, user=user)) + + # Batch the user creation in batches of 50 users + if len(activated_connections) > 50: + AccountsUsers.objects.bulk_create(activated_connections) + activated_connections = [] + + if activated_connections: + AccountsUsers.objects.bulk_create(activated_connections) + + okta_config, created = OktaSettings.objects.get_or_create(account=account) for field in ["client_id", "client_secret", "url", "enabled", "enforced"]: value = getattr(typed_input, field) diff --git a/codecov_auth/commands/owner/interactors/tests/test_save_okta_config.py b/codecov_auth/commands/owner/interactors/tests/test_save_okta_config.py index 038910acee..e9de477d2e 100644 --- a/codecov_auth/commands/owner/interactors/tests/test_save_okta_config.py +++ b/codecov_auth/commands/owner/interactors/tests/test_save_okta_config.py @@ -2,6 +2,7 @@ from asgiref.sync import async_to_sync from django.contrib.auth.models import AnonymousUser from django.test import TransactionTestCase +from shared.django_apps.codecov_auth.models import Account from codecov.commands.exceptions import Unauthenticated, Unauthorized, ValidationError from codecov_auth.models import OktaSettings @@ -18,16 +19,20 @@ class SaveOktaConfigInteractorTest(TransactionTestCase): def setUp(self): self.current_user = OwnerFactory(username="codecov-user") self.service = "github" + user1 = OwnerFactory() + user2 = OwnerFactory() self.owner = OwnerFactory( username=self.current_user.username, service=self.service, account=AccountFactory(), ) + self.owner_with_admins = OwnerFactory( username=self.current_user.username, service=self.service, admins=[self.current_user.ownerid], - account=AccountFactory(), + plan_activated_users=[user1.ownerid, user2.ownerid], + account=None, ) self.interactor = SaveOktaConfigInteractor( @@ -88,13 +93,24 @@ def test_unauthorized_error_when_user_is_not_admin(self): ) def test_create_okta_settings_when_account_does_not_exist(self): + plan_activated_users = [] + for _ in range(100): + user_owner = OwnerFactory() + plan_activated_users.append(user_owner.ownerid) + + org_with_lots_of_users = OwnerFactory( + service=self.service, + admins=[self.current_user.ownerid], + plan_activated_users=plan_activated_users, + ) + input_data = { "client_id": "some-client-id", "client_secret": "some-client-secret", "url": "https://okta.example.com", "enabled": True, "enforced": True, - "org_username": self.owner_with_admins.username, + "org_username": org_with_lots_of_users.username, } interactor = SaveOktaConfigInteractor( @@ -102,7 +118,18 @@ def test_create_okta_settings_when_account_does_not_exist(self): ) self.execute(interactor=interactor, input=input_data) - okta_config = OktaSettings.objects.get(account=self.owner_with_admins.account) + org_with_lots_of_users.refresh_from_db() + account = org_with_lots_of_users.account + + assert account.name == org_with_lots_of_users.username + assert account.plan == org_with_lots_of_users.plan + assert account.plan_seat_count == org_with_lots_of_users.plan_user_count + assert account.free_seat_count == org_with_lots_of_users.free + + assert account.users.count() == 100 + assert account.users.count() == len(org_with_lots_of_users.plan_activated_users) + + okta_config = OktaSettings.objects.get(account=org_with_lots_of_users.account) assert okta_config.client_id == input_data["client_id"] assert okta_config.client_secret == input_data["client_secret"]