GML-1660 change the importing make_connection

tigergraph · Jun 12, 2024 · 3403e3c · 3403e3c
1 parent eebf27e
commit 3403e3c
Showing 1 changed file with 55 additions and 64 deletions.
diff --git a/tests/test_jwtAuth.py b/tests/test_jwtAuth.py
@@ -1,95 +1,86 @@
+import unittest
 import requests
 import json
-import os
-from os.path import exists
-from requests.auth import HTTPBasicAuth
-
-import unittest
-from unittest.mock import patch, Mock, MagicMock
 
+from pyTigerGraphUnitTest import make_connection
 from pyTigerGraph import TigerGraphConnection
-from tests.pyTigerGraphUnitTest import make_connection
+
 
 class TestJWTTokenAuth(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.conn = make_connection(graphname="mygraph")
+        cls.conn = make_connection(graphname="Cora")
 
 
     def test_jwtauth(self):
-        # self.conn.gsql("DROP GRAPH jwttoken")
-        # self.conn.gsql("CREATE GRAPH tests()")
-        # print (self.conn.graphname)
-        # self.conn.gsql("USE GRAPH mygraph")
-        # authheader = self.conn.authHeader
-        # print (f"authheader from init conn: {authheader}")
+        print (self.conn.authHeader)
         dbversion = self.conn.getVer()
         print (f"dbversion from init conn: {dbversion}")
 
-        # if "3.9" in str(dbversion):
-        #     self._test_jwtauth_3_9()
-        # elif "4.1" in str(dbversion):
-        #     self._test_jwtauth_4_1_success()
-        #     self._test_jwtauth_4_1_fail()
-        # else:
-        #     pass
+        if "3.9" in str(dbversion):
+            self._test_jwtauth_3_9()
+        elif "4.1" in str(dbversion):
+            self._test_jwtauth_4_1_success()
+            self._test_jwtauth_4_1_fail()
+        else:
+            pass
 
 
-    # def _requestJWTToken(self):
-    #     # Define the URL
-    #     url = f"{self.conn.host}:{self.conn.gsPort}/gsqlserver/requestjwttoken"
-    #     # Define the data payload
-    #     payload = json.dumps({"lifetime": "1000000000"})
-    #     # Define the headers for the request
-    #     headers = {
-    #         'Content-Type': 'application/json'
-    #     }
-    #     # Make the POST request with basic authentication
-    #     response = requests.post(url, data=payload, headers=headers, auth=(self.conn.username, self.conn.password))
-    #     return response.json()['token']
+    def _requestJWTToken(self):
+        # Define the URL
+        url = f"{self.conn.host}:{self.conn.gsPort}/gsqlserver/requestjwttoken"
+        # Define the data payload
+        payload = json.dumps({"lifetime": "1000000000"})
+        # Define the headers for the request
+        headers = {
+            'Content-Type': 'application/json'
+        }
+        # Make the POST request with basic authentication
+        response = requests.post(url, data=payload, headers=headers, auth=(self.conn.username, self.conn.password))
+        return response.json()['token']
 
 
-    # def _test_jwtauth_3_9(self):
-    #     with self.assertRaises(RuntimeError) as context:
-    #         TigerGraphConnection(
-    #             host=self.conn.host,
-    #             jwtToken="fake.JWT.Token"
-    #         )
+    def _test_jwtauth_3_9(self):
+        with self.assertRaises(RuntimeError) as context:
+            TigerGraphConnection(
+                host=self.conn.host,
+                jwtToken="fake.JWT.Token"
+            )
 
-    #     # Verify the exception message
-    #     self.assertIn("switch to API token or username/password.", str(context.exception))
+        # Verify the exception message
+        self.assertIn("switch to API token or username/password.", str(context.exception))
 
 
-    # def _test_jwtauth_4_1_success(self):
-    #     jwt_token = self._requestJWTToken()
+    def _test_jwtauth_4_1_success(self):
+        jwt_token = self._requestJWTToken()
 
-    #     newconn = TigerGraphConnection(
-    #         host=self.conn.host,
-    #         jwtToken=jwt_token
-    #     )
+        newconn = TigerGraphConnection(
+            host=self.conn.host,
+            jwtToken=jwt_token
+        )
 
-    #     authheader = newconn.authHeader
-    #     print (f"authheader from new conn: {authheader}")
+        authheader = newconn.authHeader
+        print (f"authheader from new conn: {authheader}")
 
-    #     # restpp on port 9000
-    #     dbversion = newconn.getVer()
-    #     print (f"dbversion from new conn: {dbversion}")
-    #     self.assertIn("4.1", str(dbversion))
+        # restpp on port 9000
+        dbversion = newconn.getVer()
+        print (f"dbversion from new conn: {dbversion}")
+        self.assertIn("4.1", str(dbversion))
 
-    #     # gsql on port 14240
-    #     res = newconn._get(f"{self.conn.host}:{self.conn.gsPort}/gsqlserver/gsql/simpleauth", authMode="token", resKey=None)
-    #     self.assertIn("privileges", res)
+        # gsql on port 14240
+        res = newconn._get(f"{self.conn.host}:{self.conn.gsPort}/gsqlserver/gsql/simpleauth", authMode="token", resKey=None)
+        self.assertIn("privileges", res)
 
 
-    # def _test_jwtauth_4_1_fail(self):
-    #     with self.assertRaises(RuntimeError) as context:
-    #         TigerGraphConnection(
-    #             host=self.conn.host,
-    #             jwtToken="invalid.JWT.Token"
-    #         )
+    def _test_jwtauth_4_1_fail(self):
+        with self.assertRaises(RuntimeError) as context:
+            TigerGraphConnection(
+                host=self.conn.host,
+                jwtToken="invalid.JWT.Token"
+            )
 
-    #     # Verify the exception message
-    #     self.assertIn("Please generate new JWT token", str(context.exception))
+        # Verify the exception message
+        self.assertIn("Please generate new JWT token", str(context.exception))
 
 
 if __name__ == '__main__':