Skip to content

Commit 4c99954

Browse files
committed
Reuse fixture from conftest.py
1 parent fdb681a commit 4c99954

File tree

2 files changed

+28
-16
lines changed

2 files changed

+28
-16
lines changed

instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,8 @@ def instrument_no_content(
111111

112112
yield instrumentor
113113
os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None)
114-
instrumentor.uninstrument()
114+
if instrumentor.is_instrumented_by_opentelemetry:
115+
instrumentor.uninstrument()
115116

116117

117118
@pytest.fixture
@@ -130,7 +131,8 @@ def instrument_with_content(
130131

131132
yield instrumentor
132133
os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None)
133-
instrumentor.uninstrument()
134+
if instrumentor.is_instrumented_by_opentelemetry:
135+
instrumentor.uninstrument()
134136

135137

136138
@pytest.fixture(scope="module")
Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,19 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
117
import pytest
218
from google.cloud.aiplatform_v1.services.prediction_service import client
319
from google.cloud.aiplatform_v1beta1.services.prediction_service import (
@@ -7,31 +23,25 @@
723
from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
824

925

10-
@pytest.fixture(name="instrumentor")
11-
def fixture_instrumentor():
12-
instrumentor = VertexAIInstrumentor()
13-
instrumentor.instrument()
14-
yield instrumentor
15-
16-
if instrumentor.is_instrumented_by_opentelemetry:
17-
instrumentor.uninstrument()
18-
19-
2026
@pytest.fixture(
2127
name="client_class",
2228
params=[
2329
pytest.param(client.PredictionServiceClient, id="v1"),
2430
pytest.param(client_v1beta1.PredictionServiceClient, id="v1beta1"),
2531
],
2632
)
27-
def fixture_client_class(request):
33+
def fixture_client_class(request: pytest.FixtureRequest):
2834
return request.param
2935

3036

31-
def test_instruments(instrumentor: VertexAIInstrumentor, client_class):
37+
def test_instruments(
38+
instrument_with_content: VertexAIInstrumentor, client_class
39+
):
3240
assert hasattr(client_class.generate_content, "__wrapped__")
3341

3442

35-
def test_uninstruments(instrumentor: VertexAIInstrumentor, client_class):
36-
instrumentor.uninstrument()
43+
def test_uninstruments(
44+
instrument_with_content: VertexAIInstrumentor, client_class
45+
):
46+
instrument_with_content.uninstrument()
3747
assert not hasattr(client_class.generate_content, "__wrapped__")

0 commit comments

Comments
 (0)