diff --git a/src/scripts/create_teams.py b/src/scripts/create_teams.py index 6ff5df9..b834b41 100644 --- a/src/scripts/create_teams.py +++ b/src/scripts/create_teams.py @@ -10,7 +10,7 @@ from src.shared.clients.jit import get_existing_teams, create_teams, list_assets, add_teams_to_asset, delete_teams, \ get_jit_jwt_token from src.shared.diff_tools import get_different_items_in_lists -from src.shared.models import Asset, TeamAttributes, Organization, TeamStructure +from src.shared.models import Asset, TeamAttributes, Organization, TeamStructure, ResourceType # Load environment variables from .env file. load_dotenv() @@ -82,7 +82,8 @@ def update_assets(token, assets: List[Asset], organization): logger.info(f"Syncing team(s) {teams_to_update} to asset '{asset.asset_name}'") add_teams_to_asset(token, asset, teams_to_update) else: - if asset.tags and "team" in [t.name for t in asset.tags]: + asset_has_teams_tag = asset.tags and "team" in [t.name for t in asset.tags] + if asset_has_teams_tag: logger.info(f"Removing all teams from asset '{asset.asset_name}'") add_teams_to_asset(token, asset, teams_to_update) @@ -131,7 +132,7 @@ def get_desired_teams(assets: List[Asset], organization: Organization) -> List[s for team in organization.teams: team_resources = [] for resource in team.resources: - if resource.type == "github_repo" and resource.name in [asset.asset_name for asset in assets]: + if resource.type == ResourceType.GithubRepo and resource.name in [asset.asset_name for asset in assets]: team_resources.append(resource.name) if team_resources: desired_teams.append(team.name) @@ -189,7 +190,7 @@ def get_teams_for_assets(organization: Organization) -> Dict[str, List[str]]: asset_to_team_map = {} for team in organization.teams: for resource in team.resources: - if resource.type == "github_repo": + if resource.type == ResourceType.GithubRepo: asset_name = resource.name if asset_name in asset_to_team_map: asset_to_team_map[asset_name].append(team.name) diff --git a/src/shared/clients/github.py b/src/shared/clients/github.py index 75991bd..6adee9f 100644 --- a/src/shared/clients/github.py +++ b/src/shared/clients/github.py @@ -3,7 +3,7 @@ from github import Github from loguru import logger -from src.shared.models import TeamStructure, Resource, Organization +from src.shared.models import TeamStructure, Resource, Organization, ResourceType def get_teams_from_github_topics() -> Organization: @@ -33,11 +33,11 @@ def get_teams_from_github_topics() -> Organization: # Check if the topic already exists in the teams dictionary if topic in teams: # Add the repository to the existing team - teams[topic].resources.append(Resource(type="github_repo", name=repo_name)) + teams[topic].resources.append(Resource(type=ResourceType.GithubRepo, name=repo_name)) else: # Create a new team template for the topic team_template = TeamStructure(name=topic, members=[], - resources=[Resource(type="github_repo", name=repo_name)]) + resources=[Resource(type=ResourceType.GithubRepo, name=repo_name)]) # Add the team template to the teams dictionary teams[topic] = team_template diff --git a/src/shared/clients/jit.py b/src/shared/clients/jit.py index 503d9ab..96b8d7d 100644 --- a/src/shared/clients/jit.py +++ b/src/shared/clients/jit.py @@ -6,6 +6,7 @@ from loguru import logger from src.shared.consts import JIT_DEFAULT_API_ENDPOINT +from src.shared.env_tools import get_jit_endpoint_base_url from src.shared.models import Asset, TeamAttributes @@ -15,7 +16,7 @@ def get_jit_jwt_token() -> Optional[str]: "secret": os.getenv('JIT_CLIENT_SECRET') } - response = requests.post(f"{os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT)}/authentication/login", + response = requests.post(f"{get_jit_endpoint_base_url()}/authentication/login", json=payload) if response.status_code == 200: @@ -28,7 +29,7 @@ def get_jit_jwt_token() -> Optional[str]: def list_assets(token: str) -> List[Asset]: try: # Make a GET request to the asset API - url = f"{os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT)}/asset" + url = f"{get_jit_endpoint_base_url()}/asset" headers = { "Authorization": f"Bearer {token}" } @@ -59,7 +60,7 @@ def _handle_resoponse(response, existing_teams): try: # Make a GET request to the asset API - url = f"{os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT)}/teams?limit=100" + url = f"{get_jit_endpoint_base_url()}/teams?limit=100" headers = get_request_headers(token) response = requests.get(url, headers=headers) @@ -96,7 +97,7 @@ def delete_teams(token, team_names): break if team_id: - url = f"{os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT)}/teams/{team_id}" + url = f"{get_jit_endpoint_base_url()}/teams/{team_id}" headers = {"Authorization": f"Bearer {token}"} response = requests.delete(url, headers=headers) @@ -112,7 +113,7 @@ def delete_teams(token, team_names): def create_teams(token, teams_to_create): try: - url = f"{os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT)}/teams/" + url = f"{get_jit_endpoint_base_url()}/teams/" headers = get_request_headers(token) for team_name in teams_to_create: payload = { @@ -137,7 +138,7 @@ def get_request_headers(token): def add_teams_to_asset(token, asset: Asset, teams: List[str]): try: - url = f"{os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT)}/asset/asset/{asset.asset_id}" + url = f"{get_jit_endpoint_base_url()}/asset/asset/{asset.asset_id}" headers = get_request_headers(token) payload = { "teams": teams diff --git a/src/shared/env_tools.py b/src/shared/env_tools.py new file mode 100644 index 0000000..d27595e --- /dev/null +++ b/src/shared/env_tools.py @@ -0,0 +1,7 @@ +import os + +from src.shared.consts import JIT_DEFAULT_API_ENDPOINT + + +def get_jit_endpoint_base_url() -> str: + return os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT) diff --git a/src/shared/models.py b/src/shared/models.py index 383e1e8..1849d7a 100644 --- a/src/shared/models.py +++ b/src/shared/models.py @@ -1,10 +1,14 @@ +from enum import Enum from typing import Optional, List from pydantic import BaseModel - from src.shared.consts import MANUAL_TEAM_SOURCE +class ResourceType(str, Enum): + GithubRepo = 'github_repo' + + class TeamAttributes(BaseModel): tenant_id: str id: str diff --git a/tests/factories.py b/tests/factories.py index 4a9aa37..3bae738 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -3,7 +3,7 @@ from faker import Faker from polyfactory.factories.pydantic_factory import ModelFactory from src.shared.consts import MANUAL_TEAM_SOURCE -from src.shared.models import TeamAttributes, TeamStructure, Asset, Organization, Resource +from src.shared.models import TeamAttributes, TeamStructure, Asset, Organization, Resource, ResourceType locales = OrderedDict([ ('en-US', 1), @@ -29,7 +29,7 @@ class TeamAttributesFactory(ModelFactory): class ResourceFactory(ModelFactory): __model__ = Resource - type = "github_repo" + type = ResourceType.GithubRepo name = fake.word diff --git a/tests/shared/test_clients.py b/tests/shared/test_clients.py index 30398d9..1ce997f 100644 --- a/tests/shared/test_clients.py +++ b/tests/shared/test_clients.py @@ -1,12 +1,9 @@ -import os - import pytest - from src.shared.clients.github import get_teams_from_github_topics from src.shared.clients.jit import list_assets, get_existing_teams, create_teams, add_teams_to_asset, delete_teams, \ get_jit_jwt_token -from src.shared.consts import JIT_DEFAULT_API_ENDPOINT -from src.shared.models import TeamAttributes, Asset, Organization, TeamStructure, Resource +from src.shared.env_tools import get_jit_endpoint_base_url +from src.shared.models import TeamAttributes, Asset, Organization, TeamStructure, Resource, ResourceType class MockRepo: @@ -25,20 +22,22 @@ def get_topics(self): ([MockRepo("repo1", [])], Organization(teams=[])), ([MockRepo("repo1", ["topic1"]), MockRepo("repo2", ["topic2"])], Organization(teams=[ - TeamStructure(name="topic1", members=[], resources=[Resource(type="github_repo", name="repo1")]), - TeamStructure(name="topic2", members=[], resources=[Resource(type="github_repo", name="repo2")]) + TeamStructure(name="topic1", members=[], + resources=[Resource(type=ResourceType.GithubRepo, name="repo1")]), + TeamStructure(name="topic2", members=[], + resources=[Resource(type=ResourceType.GithubRepo, name="repo2")]) ])), ([MockRepo("repo1", ["topic1"]), MockRepo("repo2", ["topic2"]), MockRepo("repo3", ["topic2"]), MockRepo("repo4", ["topic1", "topic2"])], Organization(teams=[ TeamStructure(name="topic1", members=[], resources=[ - Resource(type="github_repo", name="repo1"), - Resource(type="github_repo", name="repo4") + Resource(type=ResourceType.GithubRepo, name="repo1"), + Resource(type=ResourceType.GithubRepo, name="repo4") ]), TeamStructure(name="topic2", members=[], resources=[ - Resource(type="github_repo", name="repo2"), - Resource(type="github_repo", name="repo3"), - Resource(type="github_repo", name="repo4") + Resource(type=ResourceType.GithubRepo, name="repo2"), + Resource(type=ResourceType.GithubRepo, name="repo3"), + Resource(type=ResourceType.GithubRepo, name="repo4") ]) ])), ] @@ -93,7 +92,7 @@ def test_get_jwt_token(status_code, expected_result, mocker): token = get_jit_jwt_token() requests_post_mock.assert_called_once_with( - f"{os.getenv('JIT_API_ENDPOINT', JIT_DEFAULT_API_ENDPOINT)}/authentication/login", + f"{get_jit_endpoint_base_url()}/authentication/login", json={"clientId": None, "secret": None} ) assert token == expected_result @@ -205,7 +204,7 @@ def test_delete_teams(mocker, status_code, existing_team_names, input_team_names expected_warning): mock_existing_teams = [ TeamAttributes(tenant_id=f"tenant{i + 1}", id=str(i + 1), created_at=f"date{i + 1}", modified_at=f"date{i + 2}", - name=team_name) + name=team_name) for i, team_name in enumerate(existing_team_names) ]