Allow tenant to be configurable

Author: FlorentinD

scripts/ci/aura_api_ci.py

@@ -1,14 +1,10 @@
 import logging
-import os
 import time
 from time import sleep
 from typing import Any, Dict, Optional
 
 import requests
 
-CLIENT_ID = os.environ["AURA_API_CLIENT_ID"]
-CLIENT_SECRET = os.environ["AURA_API_CLIENT_SECRET"]
-
 
 class AuraApiCI:
     class AuraAuthToken:
@@ -25,9 +21,11 @@ def __init__(self, json: Dict[str, Any]) -> None:
         def is_expired(self) -> bool:
             return self.expires_at >= int(time.time())
 
-    def __init__(self) -> None:
+    def __init__(self, client_id: str, client_secret: str, tenant_id: Optional[str] = None) -> None:
         self._token: Optional[AuraApiCI.AuraAuthToken] = None
         self._logger = logging.getLogger()
+        self._auth = (client_id, client_secret)
+        self._tenant_id = tenant_id
 
     def _build_header(self) -> Dict[str, str]:
         return {"Authorization": f"Bearer {self._auth_token()}", "User-agent": "neo4j-graphdatascience-ci"}
@@ -44,7 +42,7 @@ def _update_token(self) -> AuraAuthToken:
 
         self._logger.debug("Updating oauth token")
 
-        response = requests.post("https://api-staging.neo4j.io/oauth/token", data=data, auth=(CLIENT_ID, CLIENT_SECRET))
+        response = requests.post("https://api-staging.neo4j.io/oauth/token", data=data, auth=self._auth)
 
         response.raise_for_status()
 
@@ -143,6 +141,9 @@ def teardown_instance(self, db_id: str) -> None:
         response.raise_for_status()
 
     def get_tenant_id(self) -> str:
+        if self._tenant_id:
+            return self._tenant_id
+
         response = requests.get(
             "https://api-staging.neo4j.io/v1/tenants",
             headers=self._build_header(),

scripts/ci/run_targeting_aura.py

@@ -26,7 +26,10 @@ def run_notebooks(uri: str, username: str, password: str) -> None:
 
 
 def main() -> None:
-    aura_api = AuraApiCI()
+    client_id = os.environ["AURA_API_CLIENT_ID"]
+    client_secret = os.environ["AURA_API_CLIENT_SECRET"]
+    tenant_id = os.environ.get("TENANT_ID")
+    aura_api = AuraApiCI(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id)
 
     MAX_INT = 1000000
     instance_name = f"ci-build-{sys.argv[2]}" if len(sys.argv) > 1 else "ci-instance-" + str(rd.randint(0, MAX_INT))

scripts/ci/run_targeting_aura_sessions.py

@@ -11,10 +11,13 @@
 
 
 def main() -> None:
-    aura_api = AuraApiCI()
+    client_id = os.environ["AURA_API_CLIENT_ID"]
+    client_secret = os.environ["AURA_API_CLIENT_SECRET"]
+    tenant_id = os.environ.get("TENANT_ID")
+    aura_api = AuraApiCI(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id)
 
     MAX_INT = 1000000
-    instance_name = f"ci-build-{sys.argv[1]}" if len(sys.argv) > 0 else "ci-instance-" + str(rd.randint(0, MAX_INT))
+    instance_name = f"ci-build-{sys.argv[1]}" if len(sys.argv) > 1 else "ci-instance-" + str(rd.randint(0, MAX_INT))
 
     create_result = aura_api.create_instance(instance_name, memory="1GB", type="professional-db")
     instance_id = create_result["id"]