1717
1818import google .genai
1919
20+ from .auth import FakeCredentials
2021from .instrumentation_context import InstrumentationContext
2122from .otel_mocker import OTelMocker
22- from .requests_mocker import RequestsMocker
23-
24-
25- class _FakeCredentials (google .auth .credentials .AnonymousCredentials ):
26- def refresh (self , request ):
27- pass
2823
2924
3025class TestCase (unittest .TestCase ):
3126 def setUp (self ):
3227 self ._otel = OTelMocker ()
3328 self ._otel .install ()
34- self ._requests = RequestsMocker ()
35- self ._requests .install ()
3629 self ._instrumentation_context = None
3730 self ._api_key = "test-api-key"
3831 self ._project = "test-project"
3932 self ._location = "test-location"
4033 self ._client = None
4134 self ._uses_vertex = False
42- self ._credentials = _FakeCredentials ()
35+ self ._credentials = FakeCredentials ()
4336
4437 def _lazy_init (self ):
4538 self ._instrumentation_context = InstrumentationContext ()
@@ -51,17 +44,22 @@ def client(self):
5144 self ._client = self ._create_client ()
5245 return self ._client
5346
54- @property
55- def requests (self ):
56- return self ._requests
57-
5847 @property
5948 def otel (self ):
6049 return self ._otel
6150
6251 def set_use_vertex (self , use_vertex ):
6352 self ._uses_vertex = use_vertex
6453
54+ def reset_client (self ):
55+ self ._client = None
56+
57+ def reset_instrumentation (self ):
58+ if self ._instrumentation_context is None :
59+ return
60+ self ._instrumentation_context .uninstall ()
61+ self ._instrumentation_context = None
62+
6563 def _create_client (self ):
6664 self ._lazy_init ()
6765 if self ._uses_vertex :
@@ -72,10 +70,9 @@ def _create_client(self):
7270 location = self ._location ,
7371 credentials = self ._credentials ,
7472 )
75- return google .genai .Client (api_key = self ._api_key )
73+ return google .genai .Client (vertexai = False , api_key = self ._api_key )
7674
7775 def tearDown (self ):
7876 if self ._instrumentation_context is not None :
7977 self ._instrumentation_context .uninstall ()
80- self ._requests .uninstall ()
8178 self ._otel .uninstall ()
0 commit comments