add get hubs

chuongmep · Sep 10, 2024 · 13ec84b · 13ec84b
1 parent 1e60206
commit 13ec84b
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 16 deletions.
diff --git a/.gitignore b/.gitignore
@@ -404,3 +404,4 @@ database
 
 **/_trial_temp*
 *.DS_Store
+*refresh_token.json
diff --git a/APSToolkitPython/src/aps_toolkit/AECDataModel.py b/APSToolkitPython/src/aps_toolkit/AECDataModel.py
@@ -1,3 +1,5 @@
+import pandas as pd
+
 from .Token import Token
 import requests
 
@@ -26,3 +28,23 @@ def execute_query_variables(self, query, variables):
         if response.status_code != 200:
             raise Exception(f"Error: {response.content}")
         return response.json()
+
+    def get_hubs(self) -> pd.DataFrame:
+        data = {
+            "query": """
+                query GetHubs {
+                    hubs {
+                        results {
+                            id
+                            name
+                            alternativeIdentifiers{
+                            dataManagementAPIHubId
+                            }
+                        }
+                    }
+                }
+            """
+        }
+        result = self.execute_query(data)
+        hubs = result['data']['hubs']['results']
+        return pd.json_normalize(hubs)
diff --git a/APSToolkitPython/src/test/test_aec_data_model.py b/APSToolkitPython/src/test/test_aec_data_model.py
@@ -6,27 +6,25 @@
 
 class TestAECDataModel(TestCase):
     def setUp(self):
-        self.token = Auth().auth3leg()
+        # read refresh token from json
+        if not os.path.exists('refresh_token.json'):
+            with open('refresh_token.json', 'w') as f:
+                f.write('')
+        with open('refresh_token.json', 'r') as f:
+            self.refresh_token = f.read()
+        self.token = Auth.refresh_token_from_env(self.refresh_token)
+        self.refresh_token = self.token.refresh_token
+        # save to json
+        with open('refresh_token.json', 'w') as f:
+            f.write(self.refresh_token)
         self.hub_id = "urn:adsk.ace:prod.scope:207eaa13-b1e2-4e99-9f20-dc94fc599272"
         self.project_id = "urn:adsk.workspace:prod.project:f4d95147-4eef-4664-89e0-a4f01f8a7b71"
         self.group_id = "YWVjZH45enZvNHRHazl1RTI4VVc0NUsySkgzX0wyQ35JZE10SWtTU1I2eXBaTjFfLTJVd0RR"
         self.aec_data_model = AECDataModel(self.token)
 
     def test_get_hubs(self):
-        data = {
-            "query": """
-                query GetHubs {
-                    hubs {
-                        results {
-                            id
-                            name
-                        }
-                    }
-                }
-            """
-        }
-        result = self.aec_data_model.execute_query(data)
-        self.assertIsNotNone(result)
+        df_result = self.aec_data_model.get_hubs()
+        self.assertIsNotNone(df_result)
 
     def test_get_projects(self):
         data = {
@@ -188,4 +186,4 @@ def test_get_data_instance_type_category(self):
             "propertyFilter": "'property.name.category'==Walls and 'property.name.Element Context'==Instance"
         }
         result = self.aec_data_model.execute_query_variables(query, variables)
-        self.assertIsNotNone(result)
+        self.assertIsNotNone(result)