cr fixes

Author: k-shlomi

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)

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

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

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)

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

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
 
 

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)
     ]