Refactor optional test_access_token arg

bcostaaa01 · Oct 31, 2023 · 6312c65 · 6312c65
1 parent 3962132
commit 6312c65
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/hubmigrate/classes/association.py b/hubmigrate/classes/association.py
@@ -8,7 +8,7 @@ def __init__(self, config, hubspot, test_access_token=None):
         self.config = config
         self.hubspot = hubspot
         self.base_path = 'https://api.hubapi.com/crm-associations/v1/associations'
-        self.headers = {"Authorization": f"Bearer { test_access_token if test_access_token else Auth.get_token() }"}
+        self.headers = {"Authorization": f"Bearer { Auth.get_token() }"}
 
 
     def associate_records(cls, object1_id, object2_id, definition_id):
@@ -30,9 +30,7 @@ def associate_records(cls, object1_id, object2_id, definition_id):
 
         response = requests.put(url, json=payload, headers=cls.headers)
 
-        status_codes = {200, 201, 202, 204}
-
-        if response.status_code in status_codes:
+        if response.status_code == 204:
             print(f"Successfully associated records {object1_id} and {object2_id} 🎉")
         else:
             print(f"Error associating records {object1_id} and {object2_id}: {response.status_code} - {response.text} ❌")

diff --git a/hubmigrate/classes/migration.py b/hubmigrate/classes/migration.py
@@ -9,7 +9,7 @@ def __init__(self, object_type, config, hubspot, test_access_token=None):
         self.hubspot = hubspot
         self.base_path = 'https://api.hubapi.com/crm/v3/objects'
         self.path = f"{self.base_path}/{object_type}"
-        self.headers = {"Authorization": f"Bearer { test_access_token if test_access_token else Auth.get_token() }"}
+        self.headers = {"Authorization": f"Bearer { Auth.get_token() }"}
 
     status_codes = {
         200,