@@ -2039,6 +2039,7 @@ def test_list_tables(self):
2039
2039
2040
2040
def session (self , ** kwargs ):
2041
2041
from google .cloud .spanner_v1 .session import Session
2042
+
2042
2043
return Session (database = self , ** kwargs )
2043
2044
2044
2045
@property
@@ -2469,7 +2470,9 @@ def test_ctor_w_exact_staleness(self):
2469
2470
def test_from_dict (self ):
2470
2471
klass = self ._get_target_class ()
2471
2472
database = self ._make_database ()
2472
- session = database ._session_manager .get_session .return_value = self ._make_session ()
2473
+ session = (
2474
+ database ._session_manager .get_session .return_value
2475
+ ) = self ._make_session ()
2473
2476
snapshot = session .snapshot .return_value = self ._make_snapshot ()
2474
2477
api_repr = {
2475
2478
"session_id" : self .SESSION_ID ,
@@ -2504,7 +2507,9 @@ def test__get_session_already(self):
2504
2507
2505
2508
def test__get_session_new (self ):
2506
2509
database = self ._make_database ()
2507
- session = database ._session_manager .get_session .return_value = self ._make_session ()
2510
+ session = (
2511
+ database ._session_manager .get_session .return_value
2512
+ ) = self ._make_session ()
2508
2513
batch_txn = self ._make_one (database )
2509
2514
self .assertIs (batch_txn ._get_session (), session )
2510
2515
session .create .assert_called_once_with ()
@@ -3476,28 +3481,29 @@ def __init__(self, name, instance=None):
3476
3481
3477
3482
# Create a mock session manager for the new architecture
3478
3483
self ._session_manager = type ("MockSessionManager" , (), {})()
3479
-
3484
+
3480
3485
# Add get_session method to mock session manager
3481
3486
def mock_get_session (transaction_type = None ):
3482
3487
# Return a session from the pool if available
3483
- if hasattr (self ._session_manager , ' _pool' ) and self ._session_manager ._pool :
3488
+ if hasattr (self ._session_manager , " _pool" ) and self ._session_manager ._pool :
3484
3489
return self ._session_manager ._pool .get ()
3485
3490
# Otherwise create a new session
3486
3491
return _Session (database = self )
3487
-
3492
+
3488
3493
def mock_put_session (session ):
3489
3494
# Put session back into the pool if available
3490
- if hasattr (self ._session_manager , ' _pool' ) and self ._session_manager ._pool :
3495
+ if hasattr (self ._session_manager , " _pool" ) and self ._session_manager ._pool :
3491
3496
try :
3492
3497
self ._session_manager ._pool .put (session )
3493
3498
except :
3494
3499
pass # Ignore errors when putting back sessions
3495
-
3500
+
3496
3501
self ._session_manager .get_session = mock_get_session
3497
3502
self ._session_manager .put_session = mock_put_session
3498
3503
3499
3504
def session (self , ** kwargs ):
3500
3505
from google .cloud .spanner_v1 .session import Session
3506
+
3501
3507
return Session (database = self , ** kwargs )
3502
3508
3503
3509
@property
0 commit comments