Fix tests

Author: soerenreichardt

graphdatascience/query_runner/arrow_authentication.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 
 
 class ArrowAuthentication(ABC):
@@ -8,10 +9,10 @@ def auth_pair(self) -> tuple[str, str]:
         pass
 
 
+@dataclass
 class UsernamePasswordAuthentication(ArrowAuthentication):
-    def __init__(self, username: str, password: str):
-        self._username = username
-        self._password = password
+    username: str
+    password: str
 
     def auth_pair(self) -> tuple[str, str]:
-        return self._username, self._password
+        return self.username, self.password

graphdatascience/tests/unit/test_aura_api.py

@@ -807,14 +807,14 @@ def test_auth_token(requests_mock: Mocker) -> None:
         json={"access_token": "very_short_token", "expires_in": 0, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth_pair._auth_token() == "very_short_token"  # type: ignore
+    assert api._request_session.auth._auth_token() == "very_short_token"  # type: ignore
 
     requests_mock.post(
         "https://api.neo4j.io/oauth/token",
         json={"access_token": "longer_token", "expires_in": 3600, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth_pair._auth_token() == "longer_token"  # type: ignore
+    assert api._request_session.auth._auth_token() == "longer_token"  # type: ignore
 
 
 def test_auth_token_reused(requests_mock: Mocker) -> None:
@@ -825,15 +825,15 @@ def test_auth_token_reused(requests_mock: Mocker) -> None:
         json={"access_token": "one_token", "expires_in": 3600, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
 
     requests_mock.post(
         "https://api.neo4j.io/oauth/token",
         json={"access_token": "new_token", "expires_in": 3600, "token_type": "Bearer"},
     )
 
     # no new token requested
-    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
 
 
 def test_auth_token_use_short_token(requests_mock: Mocker) -> None:
@@ -844,8 +844,8 @@ def test_auth_token_use_short_token(requests_mock: Mocker) -> None:
         json={"access_token": "one_token", "expires_in": 10, "token_type": "Bearer"},
     )
 
-    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
-    assert api._request_session.auth_pair._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
+    assert api._request_session.auth._auth_token() == "one_token"  # type: ignore
 
 
 def test_derive_tenant(requests_mock: Mocker) -> None:

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -20,6 +20,7 @@
     TimeParser,
     WaitResult,
 )
+from graphdatascience.session.aura_api_token_authentication import AuraApiTokenAuthentication
 from graphdatascience.session.cloud_location import CloudLocation
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.session.dedicated_sessions import DedicatedSessions
@@ -345,6 +346,9 @@ def test_create_attached_session(mocker: MockerFixture, aura_api: AuraApi) -> No
         ttl=ttl,
     )
 
+    arrow_authentication = gds_parameters["arrow_authentication"]  # type: ignore
+    del gds_parameters["arrow_authentication"]
+
     assert gds_parameters == {  # type: ignore
         "db_runner": {
             "endpoint": "neo4j+s://ffff0.databases.neo4j.io",
@@ -353,12 +357,14 @@ def test_create_attached_session(mocker: MockerFixture, aura_api: AuraApi) -> No
             "database": None,
             "show_progress": False,
         },
-        "session_connection": DbmsConnectionInfo(
+        "session_bolt_connection_info": DbmsConnectionInfo(
             uri="neo4j+s://foo.bar", username="client-id", password="client_secret"
         ),
         "session_id": "ffff0-ffff1",
     }
 
+    assert isinstance(arrow_authentication, AuraApiTokenAuthentication)
+
     assert len(sessions.list()) == 1
     actual_session = sessions.list()[0]
 
@@ -383,6 +389,9 @@ def test_create_standalone_session(mocker: MockerFixture, aura_api: AuraApi) ->
         ttl=ttl,
     )
 
+    arrow_authentication = gds_credentials["arrow_authentication"]  # type: ignore
+    del gds_credentials["arrow_authentication"]
+
     assert gds_credentials == {  # type: ignore
         "db_runner": {
             "endpoint": "neo4j+s://foo.bar",
@@ -391,12 +400,14 @@ def test_create_standalone_session(mocker: MockerFixture, aura_api: AuraApi) ->
             "database": None,
             "show_progress": False,
         },
-        "session_connection": DbmsConnectionInfo(
+        "session_bolt_connection_info": DbmsConnectionInfo(
             uri="neo4j+s://foo.bar", username="client-id", password="client_secret"
         ),
         "session_id": "None-ffff0",
     }
 
+    assert isinstance(arrow_authentication, AuraApiTokenAuthentication)
+
     assert len(sessions.list()) == 1
     actual_session = sessions.list()[0]
 
@@ -424,6 +435,10 @@ def test_get_or_create(mocker: MockerFixture, aura_api: AuraApi) -> None:
         DbmsConnectionInfo("neo4j+s://ffff0.databases.neo4j.io", "dbuser", "db_pw"),
     )
 
+    arrow_authentication = gds_args1["arrow_authentication"]  # type: ignore
+    del gds_args1["arrow_authentication"]
+    del gds_args2["arrow_authentication"]
+
     assert gds_args1 == {  # type: ignore
         "db_runner": {
             "endpoint": "neo4j+s://ffff0.databases.neo4j.io",
@@ -432,13 +447,15 @@ def test_get_or_create(mocker: MockerFixture, aura_api: AuraApi) -> None:
             "database": None,
             "show_progress": False,
         },
-        "session_connection": DbmsConnectionInfo(
+        "session_bolt_connection_info": DbmsConnectionInfo(
             uri="neo4j+s://foo.bar", username="client-id", password="client_secret"
         ),
         "session_id": "ffff0-ffff1",
     }
     assert gds_args1 == gds_args2
 
+    assert isinstance(arrow_authentication, AuraApiTokenAuthentication)
+
     assert [i.name for i in sessions.list()] == ["my-session"]