File tree

2 files changed

+28
-0
lines changed

2 files changed

+28
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -3423,6 +3423,12 @@ def schema_to_json(self, schema_list, destination):
34233423
with open(destination, mode="w") as file_obj:
34243424
return self._schema_to_json_file_object(json_schema_list, file_obj)
34253425

3426+
def __enter__(self):
3427+
return self
3428+
3429+
def __exit__(self, exc_type, exc_value, traceback):
3430+
self.close()
3431+
34263432

34273433
# pylint: disable=unused-argument
34283434
def _item_to_project(iterator, resource):
Original file line numberDiff line numberDiff line change
@@ -7218,6 +7218,28 @@ def test_list_rows_error(self):
72187218
with self.assertRaises(TypeError):
72197219
client.list_rows(1)
72207220

7221+
def test_context_manager_enter_returns_itself(self):
7222+
creds = _make_credentials()
7223+
http = object()
7224+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
7225+
7226+
with mock..object(client, "close"), client as context_var:
7227+
pass
7228+
7229+
self.assertIs(client, context_var)
7230+
7231+
def test_context_manager_exit_closes_client(self):
7232+
creds = _make_credentials()
7233+
http = object()
7234+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
7235+
7236+
fake_close = mock.Mock()
7237+
with mock..object(client, "close", fake_close):
7238+
with client:
7239+
pass
7240+
7241+
fake_close.assert_called_once()
7242+
72217243

72227244
class Test_make_job_id(unittest.TestCase):
72237245
def _call_fut(self, job_id, prefix=None):

0 commit comments

Comments
 (0)