Skip to content

Commit 317e285

Browse files
aabmassaabmass
committed
Use pytest pytest-vcr for google-genai tests
1 parent 53ff952 commit 317e285

File tree

8 files changed

+688
-96
lines changed

8 files changed

+688
-96
lines changed

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common/otel_assertions.py‎

Lines changed: 19 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -13,53 +13,15 @@
1313
# limitations under the License.
1414

1515

16-
import opentelemetry._events
17-
import opentelemetry._logs._internal
18-
import opentelemetry.metrics._internal
19-
import opentelemetry.trace
20-
from opentelemetry._events import (
21-
get_event_logger_provider,
22-
set_event_logger_provider,
23-
)
24-
from opentelemetry._logs import get_logger_provider, set_logger_provider
25-
from opentelemetry.metrics import get_meter_provider, set_meter_provider
26-
from opentelemetry.sdk._events import EventLoggerProvider
27-
from opentelemetry.sdk._logs import LoggerProvider
2816
from opentelemetry.sdk._logs.export import (
2917
InMemoryLogExporter,
30-
SimpleLogRecordProcessor,
3118
)
32-
from opentelemetry.sdk.metrics import MeterProvider
33-
from opentelemetry.sdk.metrics._internal.export import InMemoryMetricReader
34-
from opentelemetry.sdk.trace import TracerProvider
35-
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
19+
from opentelemetry.sdk.metrics.export import (
20+
InMemoryMetricReader,
21+
)
3622
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
3723
InMemorySpanExporter,
3824
)
39-
from opentelemetry.trace import get_tracer_provider, set_tracer_provider
40-
from opentelemetry.util._once import Once
41-
42-
43-
def _bypass_otel_once():
44-
opentelemetry.trace._TRACER_PROVIDER_SET_ONCE = Once()
45-
opentelemetry._logs._internal._LOGGER_PROVIDER_SET_ONCE = Once()
46-
opentelemetry._events._EVENT_LOGGER_PROVIDER_SET_ONCE = Once()
47-
opentelemetry.metrics._internal._METER_PROVIDER_SET_ONCE = Once()
48-
49-
50-
class OTelProviderSnapshot:
51-
def __init__(self):
52-
self._tracer_provider = get_tracer_provider()
53-
self._logger_provider = get_logger_provider()
54-
self._event_logger_provider = get_event_logger_provider()
55-
self._meter_provider = get_meter_provider()
56-
57-
def restore(self):
58-
_bypass_otel_once()
59-
set_tracer_provider(self._tracer_provider)
60-
set_logger_provider(self._logger_provider)
61-
set_event_logger_provider(self._event_logger_provider)
62-
set_meter_provider(self._meter_provider)
6325

6426

6527
class _LogWrapper:
@@ -113,26 +75,20 @@ def data(self):
11375
return self._metric.data
11476

11577

116-
class OTelMocker:
117-
def __init__(self):
118-
self._snapshot = None
119-
self._logs = InMemoryLogExporter()
120-
self._traces = InMemorySpanExporter()
121-
self._metrics = InMemoryMetricReader()
78+
class OTelAssertions:
79+
def __init__(
80+
self,
81+
logs_exporter: InMemoryLogExporter,
82+
span_exporter: InMemorySpanExporter,
83+
metric_reader: InMemoryMetricReader,
84+
):
85+
self._logs = logs_exporter
86+
self._traces = span_exporter
87+
self._metrics = metric_reader
12288
self._spans = []
12389
self._finished_logs = []
12490
self._metrics_data = []
12591

126-
def install(self):
127-
self._snapshot = OTelProviderSnapshot()
128-
_bypass_otel_once()
129-
self._install_logs()
130-
self._install_metrics()
131-
self._install_traces()
132-
133-
def uninstall(self):
134-
self._snapshot.restore()
135-
13692
def get_finished_logs(self):
13793
for log_data in self._logs.get_finished_logs():
13894
self._finished_logs.append(_LogWrapper(log_data))
@@ -166,9 +122,9 @@ def get_span_named(self, name):
166122
def assert_has_span_named(self, name):
167123
span = self.get_span_named(name)
168124
finished_spans = [span.name for span in self.get_finished_spans()]
169-
assert (
170-
span is not None
171-
), f'Could not find span named "{name}"; finished spans: {finished_spans}'
125+
assert span is not None, (
126+
f'Could not find span named "{name}"; finished spans: {finished_spans}'
127+
)
172128

173129
def get_event_named(self, event_name):
174130
for event in self.get_finished_logs():
@@ -182,9 +138,9 @@ def get_event_named(self, event_name):
182138
def assert_has_event_named(self, name):
183139
event = self.get_event_named(name)
184140
finished_logs = self.get_finished_logs()
185-
assert (
186-
event is not None
187-
), f'Could not find event named "{name}"; finished logs: {finished_logs}'
141+
assert event is not None, (
142+
f'Could not find event named "{name}"; finished logs: {finished_logs}'
143+
)
188144

189145
def assert_does_not_have_event_named(self, name):
190146
event = self.get_event_named(name)
@@ -200,19 +156,3 @@ def get_metrics_data_named(self, name):
200156
def assert_has_metrics_data_named(self, name):
201157
data = self.get_metrics_data_named(name)
202158
assert len(data) > 0
203-
204-
def _install_logs(self):
205-
provider = LoggerProvider()
206-
provider.add_log_record_processor(SimpleLogRecordProcessor(self._logs))
207-
set_logger_provider(provider)
208-
event_provider = EventLoggerProvider(logger_provider=provider)
209-
set_event_logger_provider(event_provider)
210-
211-
def _install_metrics(self):
212-
provider = MeterProvider(metric_readers=[self._metrics])
213-
set_meter_provider(provider)
214-
215-
def _install_traces(self):
216-
provider = TracerProvider()
217-
provider.add_span_processor(SimpleSpanProcessor(self._traces))
218-
set_tracer_provider(provider)

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/conftest.py‎

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,19 @@
77

88
import pytest
99
import vertexai
10+
import google.genai
11+
import pytest
1012
import yaml
1113
from google.auth.credentials import AnonymousCredentials
1214
from vcr import VCR
1315
from vcr.record_mode import RecordMode
1416
from vcr.request import Request
1517

16-
from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
17-
from opentelemetry.instrumentation.vertexai.utils import (
18-
OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT,
18+
from opentelemetry.instrumentation.google_genai import (
19+
GoogleGenAiSdkInstrumentor,
20+
)
21+
from opentelemetry.instrumentation.google_genai.flags import (
22+
_CONTENT_RECORDING_ENV_VAR,
1923
)
2024
from opentelemetry.sdk._events import EventLoggerProvider
2125
from opentelemetry.sdk._logs import LoggerProvider
@@ -35,7 +39,10 @@
3539
InMemorySpanExporter,
3640
)
3741

42+
from .common.otel_assertions import OTelAssertions
43+
3844
FAKE_PROJECT = "fake-project"
45+
FAKE_LOCATION = "us-central1"
3946

4047

4148
@pytest.fixture(scope="function", name="span_exporter")
@@ -79,8 +86,16 @@ def fixture_meter_provider(metric_reader):
7986
)
8087

8188

82-
@pytest.fixture(autouse=True)
83-
def vertexai_init(vcr: VCR) -> None:
89+
@pytest.fixture
90+
def otel_assertions(log_exporter, span_exporter, metric_reader):
91+
return OTelAssertions(
92+
logs_exporter=log_exporter,
93+
span_exporter=span_exporter,
94+
metric_reader=metric_reader,
95+
)
96+
97+
98+
def _genai_client(vertexai: bool, vcr: VCR) -> google.genai.Client:
8499
# When not recording (in CI), don't do any auth. That prevents trying to read application
85100
# default credentials from the filesystem or metadata server and oauth token exchange. This
86101
# is not the interesting part of our instrumentation to test.
@@ -89,47 +104,55 @@ def vertexai_init(vcr: VCR) -> None:
89104
if vcr.record_mode == RecordMode.NONE:
90105
credentials = AnonymousCredentials()
91106
project = FAKE_PROJECT
92-
vertexai.init(
93-
api_transport="rest", credentials=credentials, project=project
107+
return google.genai.Client(
108+
vertexai=vertexai, project=project, credentials=credentials
94109
)
95110

96111

112+
@pytest.fixture
113+
def genai_client(vcr: VCR) -> google.genai.Client:
114+
return _genai_client(vertexai=True, vcr=vcr)
115+
116+
117+
@pytest.fixture
118+
def genai_client_gemini(
119+
vcr: VCR, request: pytest.FixtureRequest
120+
) -> google.genai.Client:
121+
return _genai_client(vertexai=False, vcr=vcr)
122+
123+
97124
@pytest.fixture
98125
def instrument_no_content(
99126
tracer_provider, event_logger_provider, meter_provider
100127
):
101-
os.environ.update(
102-
{OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "False"}
103-
)
128+
os.environ.update({_CONTENT_RECORDING_ENV_VAR: "False"})
104129

105-
instrumentor = VertexAIInstrumentor()
130+
instrumentor = GoogleGenAiSdkInstrumentor()
106131
instrumentor.instrument(
107132
tracer_provider=tracer_provider,
108133
event_logger_provider=event_logger_provider,
109134
meter_provider=meter_provider,
110135
)
111136

112137
yield instrumentor
113-
os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None)
138+
os.environ.pop(_CONTENT_RECORDING_ENV_VAR, None)
114139
instrumentor.uninstrument()
115140

116141

117142
@pytest.fixture
118143
def instrument_with_content(
119144
tracer_provider, event_logger_provider, meter_provider
120145
):
121-
os.environ.update(
122-
{OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "True"}
123-
)
124-
instrumentor = VertexAIInstrumentor()
146+
os.environ.update({_CONTENT_RECORDING_ENV_VAR: "True"})
147+
instrumentor = GoogleGenAiSdkInstrumentor()
125148
instrumentor.instrument(
126149
tracer_provider=tracer_provider,
127150
event_logger_provider=event_logger_provider,
128151
meter_provider=meter_provider,
129152
)
130153

131154
yield instrumentor
132-
os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None)
155+
os.environ.pop(_CONTENT_RECORDING_ENV_VAR, None)
133156
instrumentor.uninstrument()
134157

135158

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/cassettes/test_genai_system_gemini.yaml‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
interactions:
2+
- request:
3+
body: |-
4+
{
5+
"contents": [
6+
{
7+
"parts": [
8+
{
9+
"text": "Does this work?"
10+
}
11+
],
12+
"role": "user"
13+
}
14+
]
15+
}
16+
headers:
17+
Accept:
18+
- '*/*'
19+
Accept-Encoding:
20+
- gzip, deflate
21+
Connection:
22+
- keep-alive
23+
Content-Length:
24+
- '72'
25+
Content-Type:
26+
- application/json
27+
user-agent:
28+
- google-genai-sdk/1.0.0 gl-python/3.13.1
29+
method: POST
30+
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
31+
response:
32+
body:
33+
string: |-
34+
{
35+
"candidates": [
36+
{
37+
"content": {
38+
"parts": [
39+
{
40+
"text": "Please provide me with the code or process you'd like me to evaluate. I need the code or a description of what you want to do to tell you if it works. \n\nFor example, you could say:\n\n* \"Does this Python code work to calculate the factorial of a number: `def factorial(n): if n == 0: return 1 else: return n * factorial(n-1)`?\"\n* \"Does this method of starting a car work: put the key in the ignition, turn it to the 'start' position?\"\n* \"Does this JavaScript code work to validate an email address: `function validateEmail(email) { ... }`\"\n\nThe more information you give me, the better I can help!\n"
41+
}
42+
],
43+
"role": "model"
44+
},
45+
"finishReason": "STOP",
46+
"avgLogprobs": -0.4348284892546825
47+
}
48+
],
49+
"usageMetadata": {
50+
"promptTokenCount": 4,
51+
"candidatesTokenCount": 156,
52+
"totalTokenCount": 160,
53+
"promptTokensDetails": [
54+
{
55+
"modality": "TEXT",
56+
"tokenCount": 4
57+
}
58+
],
59+
"candidatesTokensDetails": [
60+
{
61+
"modality": "TEXT",
62+
"tokenCount": 156
63+
}
64+
]
65+
},
66+
"modelVersion": "gemini-2.0-flash"
67+
}
68+
headers:
69+
Content-Type:
70+
- application/json; charset=UTF-8
71+
Transfer-Encoding:
72+
- chunked
73+
Vary:
74+
- Origin
75+
- X-Origin
76+
- Referer
77+
content-length:
78+
- '1215'
79+
status:
80+
code: 200
81+
message: OK
82+
version: 1

0 commit comments

Comments
 (0)