lint fixes and docs tweaks

Signed-off-by: Daniel Brauer <[email protected]>
cartography-cncf · Oct 31, 2024 · dcc5bef · dcc5bef
1 parent 213c2fc
commit dcc5bef
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 81 deletions.
diff --git a/cartography/intel/github/users.py b/cartography/intel/github/users.py
@@ -1,12 +1,12 @@
 import logging
+from copy import deepcopy
 from typing import Any
 from typing import Dict
 from typing import List
 from typing import Tuple
 
 import neo4j
 
-from copy import deepcopy
 from cartography.intel.github.util import fetch_all
 from cartography.stats import get_stats_client
 from cartography.util import merge_module_sync_metadata
@@ -72,6 +72,7 @@
     }
     """
 
+
 @timeit
 def get_users(token: str, api_url: str, organization: str) -> Tuple[List[Dict], Dict]:
     """
@@ -94,6 +95,7 @@ def get_users(token: str, api_url: str, organization: str) -> Tuple[List[Dict],
     )
     return users.edges, org
 
+
 def _get_enterprise_owners_raw(token: str, api_url: str, organization: str) -> Tuple[List[Dict], Dict]:
     """
     Function broken out for testing purposes.  See 'get_enterprise_owners' for docs.
@@ -107,6 +109,7 @@ def _get_enterprise_owners_raw(token: str, api_url: str, organization: str) -> T
     )
     return owners.edges, org
 
+
 @timeit
 def get_enterprise_owners(token: str, api_url: str, organization: str) -> Tuple[List[Dict], List[Dict], Dict]:
     """
@@ -140,9 +143,12 @@ def _mark_users_as_enterprise_owners(
 ) -> list[Dict]:
     """
     :param user_data: A list of dicts representing users - see tests.data.github.users.GITHUB_USER_DATA for shape.
-    :param user_org_data: A dict representing the organization for the user_data - see tests.data.github.users.GITHUB_ORG_DATA for shape.
-    :param affiliated_owner_data: A list of dicts representing affiliated enterprise owners (owners who are also users in the org) - see tests.data.github.users.GITHUB_ENTERPRISE_OWNER_DATA for shape.
-    :param owner_org_data: A dict representing the organization for the owner data - see tests.data.github.users.GITHUB_ORG_DATA for shape.
+    :param user_org_data: A dict representing the organization for the user_data
+        see tests.data.github.users.GITHUB_ORG_DATA for shape.
+    :param affiliated_owner_data: A list of dicts representing affiliated enterprise owners
+        (owners who are also users in the org) - see tests.data.github.users.GITHUB_ENTERPRISE_OWNER_DATA for shape.
+    :param owner_org_data: A dict representing the organization for the owner data
+        see tests.data.github.users.GITHUB_ORG_DATA for shape.
     :return: A new list of user_data dicts updated with a new property, isEnterpriseOwner
     """
 
@@ -200,6 +206,7 @@ def load_organization_users(
         UpdateTag=update_tag,
     )
 
+
 @timeit
 def load_unaffiliated_owners(
     neo4j_session: neo4j.Session, owner_data: List[Dict], org_data: Dict,
@@ -247,6 +254,7 @@ def load_unaffiliated_owners(
         UpdateTag=update_tag,
     )
 
+
 @timeit
 def sync(
         neo4j_session: neo4j.Session,
@@ -257,10 +265,19 @@ def sync(
 ) -> None:
     logger.info("Syncing GitHub users")
     user_data, user_org_data = get_users(github_api_key, github_url, organization)
-    affiliated_owner_data, unaffiliated_owner_data, owner_org_data = get_enterprise_owners(github_api_key, github_url, organization)
-    processed_user_data = _mark_users_as_enterprise_owners(user_data, user_org_data, affiliated_owner_data, owner_org_data)
+    affiliated_owner_data, unaffiliated_owner_data, owner_org_data = get_enterprise_owners(
+        github_api_key,
+        github_url, organization,
+    )
+    processed_user_data = _mark_users_as_enterprise_owners(
+        user_data, user_org_data,
+        affiliated_owner_data, owner_org_data,
+    )
     load_organization_users(neo4j_session, processed_user_data, user_org_data, common_job_parameters['UPDATE_TAG'])
-    load_unaffiliated_owners(neo4j_session, unaffiliated_owner_data, owner_org_data, common_job_parameters['UPDATE_TAG'])
+    load_unaffiliated_owners(
+        neo4j_session, unaffiliated_owner_data,
+        owner_org_data, common_job_parameters['UPDATE_TAG'],
+    )
     run_cleanup_job('github_users_cleanup.json', neo4j_session, common_job_parameters)
     merge_module_sync_metadata(
         neo4j_session,

diff --git a/tests/data/github/users.py b/tests/data/github/users.py
@@ -4,81 +4,85 @@
 }
 
 
-GITHUB_USER_DATA = ([
-    {
-        'hasTwoFactorEnabled': None,
-        'node': {
-            'url': 'https://example.com/hjsimpson',
-            'login': 'hjsimpson',
-            'name': 'Homer Simpson',
-            'isSiteAdmin': False,
-            'email': '[email protected]',
-            'company': 'Springfield Nuclear Power Plant',
+GITHUB_USER_DATA = (
+    [
+        {
+            'hasTwoFactorEnabled': None,
+            'node': {
+                'url': 'https://example.com/hjsimpson',
+                'login': 'hjsimpson',
+                'name': 'Homer Simpson',
+                'isSiteAdmin': False,
+                'email': '[email protected]',
+                'company': 'Springfield Nuclear Power Plant',
+            },
+            'role': 'MEMBER',
+        }, {
+            'hasTwoFactorEnabled': None,
+            'node': {
+                'url': 'https://example.com/lmsimpson',
+                'login': 'lmsimpson',
+                'name': 'Lisa Simpson',
+                'isSiteAdmin': False,
+                'email': '[email protected]',
+                'company': 'Simpson Residence',
+            },
+            'role': 'MEMBER',
+        }, {
+            'hasTwoFactorEnabled': True,
+            'node': {
+                'url': 'https://example.com/mbsimpson',
+                'login': 'mbsimpson',
+                'name': 'Marge Simpson',
+                'isSiteAdmin': False,
+                'email': '[email protected]',
+                'company': 'Simpson Residence',
+            },
+            'role': 'ADMIN',
         },
-        'role': 'MEMBER',
-    }, {
-        'hasTwoFactorEnabled': None,
-        'node': {
-            'url': 'https://example.com/lmsimpson',
-            'login': 'lmsimpson',
-            'name': 'Lisa Simpson',
-            'isSiteAdmin': False,
-            'email': '[email protected]',
-            'company': 'Simpson Residence',
-        },
-        'role': 'MEMBER',
-    }, {
-        'hasTwoFactorEnabled': True,
-        'node': {
-            'url': 'https://example.com/mbsimpson',
-            'login': 'mbsimpson',
-            'name': 'Marge Simpson',
-            'isSiteAdmin': False,
-            'email': '[email protected]',
-            'company': 'Simpson Residence',
-        },
-        'role': 'ADMIN',
-    }],
-    GITHUB_ORG_DATA
+    ],
+    GITHUB_ORG_DATA,
 )
 
 # Subtle differences between owner data and user data:
 # 1. owner data does not include a `hasTwoFactorEnabled` field (it in unavailable in the GraphQL query for these owners)
 # 2. an `organizationRole` field instead of a `role` field.  In owner data, membership within an org is not assumed, so
 #    there is an 'UNAFFILIATED' value for owners of an org who are not also members of it.  (Otherwise the 'OWNER'
-#    organizationRole matches the 'ADMIN' role in the user data, and the 'DIRECT_MEMBER' organizationRole matches the 'MEMBER' role.)
-GITHUB_ENTERPRISE_OWNER_DATA = ([
-    {
-        'node': {
-            'url': 'https://example.com/kbroflovski',
-            'login': 'kbroflovski',
-            'name': 'Kyle Broflovski',
-            'isSiteAdmin': False,
-            'email': '[email protected]',
-            'company': 'South Park Elementary',
-        },
-        'organizationRole': 'UNAFFILIATED',
-    }, {
-        'node': {
-            'url': 'https://example.com/mbsimpson',
-            'login': 'mbsimpson',
-            'name': 'Marge Simpson',
-            'isSiteAdmin': False,
-            'email': '[email protected]',
-            'company': 'Simpson Residence',
+#    organizationRole matches the 'ADMIN' role in the user data, and the 'DIRECT_MEMBER' organizationRole matches
+#    the 'MEMBER' role.)
+GITHUB_ENTERPRISE_OWNER_DATA = (
+    [
+        {
+            'node': {
+                'url': 'https://example.com/kbroflovski',
+                'login': 'kbroflovski',
+                'name': 'Kyle Broflovski',
+                'isSiteAdmin': False,
+                'email': '[email protected]',
+                'company': 'South Park Elementary',
+            },
+            'organizationRole': 'UNAFFILIATED',
+        }, {
+            'node': {
+                'url': 'https://example.com/mbsimpson',
+                'login': 'mbsimpson',
+                'name': 'Marge Simpson',
+                'isSiteAdmin': False,
+                'email': '[email protected]',
+                'company': 'Simpson Residence',
+            },
+            'organizationRole': 'OWNER',
+        }, {
+            'node': {
+                'url': 'https://example.com/lmsimpson',
+                'login': 'lmsimpson',
+                'name': 'Lisa Simpson',
+                'isSiteAdmin': False,
+                'email': '[email protected]',
+                'company': 'Simpson Residence',
+            },
+            'organizationRole': 'DIRECT_MEMBER',
         },
-        'organizationRole': 'OWNER',
-    }, {
-        'node': {
-            'url': 'https://example.com/lmsimpson',
-            'login': 'lmsimpson',
-            'name': 'Lisa Simpson',
-            'isSiteAdmin': False,
-            'email': '[email protected]',
-            'company': 'Simpson Residence',
-        },
-        'organizationRole': 'DIRECT_MEMBER',
-    }],
-    GITHUB_ORG_DATA
+    ],
+    GITHUB_ORG_DATA,
 )
-
diff --git a/tests/integration/cartography/intel/github/test_users.py b/tests/integration/cartography/intel/github/test_users.py
@@ -1,8 +1,9 @@
 from unittest.mock import patch
+
 import cartography.intel.github.users
-from tests.data.github.users import GITHUB_USER_DATA
-from tests.data.github.users import GITHUB_ORG_DATA
 from tests.data.github.users import GITHUB_ENTERPRISE_OWNER_DATA
+from tests.data.github.users import GITHUB_ORG_DATA
+from tests.data.github.users import GITHUB_USER_DATA
 
 TEST_UPDATE_TAG = 123456789
 TEST_JOB_PARAMS = {'UPDATE_TAG': TEST_UPDATE_TAG}
@@ -23,7 +24,8 @@ def test_sync(mock_owners, mock_users, neo4j_session):
         TEST_JOB_PARAMS,
         FAKE_API_KEY,
         TEST_GITHUB_URL,
-        TEST_GITHUB_ORG)
+        TEST_GITHUB_ORG,
+    )
 
     # Assert
 
@@ -103,7 +105,6 @@ def test_sync(mock_owners, mock_users, neo4j_session):
     assert actual_nodes == expected_nodes
 
     # Ensure hasTwoFactorEnabled has not been improperly overwritten for enterprise owners
-    # Ensure enterprise owners are identified
     nodes = neo4j_session.run(
         """
         MATCH (g:GitHubUser) RETURN g.id, g.has_2fa_enabled
@@ -121,4 +122,4 @@ def test_sync(mock_owners, mock_users, neo4j_session):
             n['g.has_2fa_enabled'],
         ) for n in nodes
     }
-    assert actual_nodes == expected_nodes
+    assert actual_nodes == expected_nodes