|
| 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 | + |
1 | 17 | import pytest |
2 | 18 | from google.cloud.aiplatform_v1.services.prediction_service import client |
3 | 19 | from google.cloud.aiplatform_v1beta1.services.prediction_service import ( |
|
7 | 23 | from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor |
8 | 24 |
|
9 | 25 |
|
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 | | - |
20 | 26 | @pytest.fixture( |
21 | 27 | name="client_class", |
22 | 28 | params=[ |
23 | 29 | pytest.param(client.PredictionServiceClient, id="v1"), |
24 | 30 | pytest.param(client_v1beta1.PredictionServiceClient, id="v1beta1"), |
25 | 31 | ], |
26 | 32 | ) |
27 | | -def fixture_client_class(request): |
| 33 | +def fixture_client_class(request: pytest.FixtureRequest): |
28 | 34 | return request.param |
29 | 35 |
|
30 | 36 |
|
31 | | -def test_instruments(instrumentor: VertexAIInstrumentor, client_class): |
| 37 | +def test_instruments( |
| 38 | + instrument_with_content: VertexAIInstrumentor, client_class |
| 39 | +): |
32 | 40 | assert hasattr(client_class.generate_content, "__wrapped__") |
33 | 41 |
|
34 | 42 |
|
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() |
37 | 47 | assert not hasattr(client_class.generate_content, "__wrapped__") |
0 commit comments