This repository was archived by the owner on Mar 13, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 35
Expand file tree
/
Copy path_helpers.py
More file actions
105 lines (79 loc) · 2.8 KB
/
_helpers.py
File metadata and controls
105 lines (79 loc) · 2.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
# Copyright 2021 Google LLC
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file or at
# https://developers.google.com/open-source/licenses/bsd
import configparser
import mock
import os
from sqlalchemy.testing import fixtures
try:
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
InMemorySpanExporter,
)
from opentelemetry.trace.status import StatusCode
trace.set_tracer_provider(TracerProvider())
HAS_OPENTELEMETRY_INSTALLED = True
except ImportError:
HAS_OPENTELEMETRY_INSTALLED = False
StatusCode = mock.Mock()
_TEST_OT_EXPORTER = None
_TEST_OT_PROVIDER_INITIALIZED = False
PROJECT = os.getenv(
"GOOGLE_CLOUD_PROJECT",
os.getenv("PROJECT_ID", "emulator-test-project"),
)
DB_URL = (
f"spanner:///projects/{PROJECT}/instances/"
"sqlalchemy-dialect-test/databases/compliance-test"
)
def get_db_url():
config = configparser.ConfigParser()
if os.path.exists("test.cfg"):
config.read("test.cfg")
else:
config.read("setup.cfg")
return config.get("db", "default", fallback=DB_URL)
def get_test_ot_exporter():
global _TEST_OT_EXPORTER
if _TEST_OT_EXPORTER is None:
_TEST_OT_EXPORTER = InMemorySpanExporter()
return _TEST_OT_EXPORTER
def use_test_ot_exporter():
global _TEST_OT_PROVIDER_INITIALIZED
if _TEST_OT_PROVIDER_INITIALIZED:
return
provider = trace.get_tracer_provider()
if not hasattr(provider, "add_span_processor"):
return
provider.add_span_processor(SimpleSpanProcessor(get_test_ot_exporter()))
_TEST_OT_PROVIDER_INITIALIZED = True
def get_project():
return PROJECT
class OpenTelemetryBase(fixtures.TestBase):
@classmethod
def setup_class(cls):
if HAS_OPENTELEMETRY_INSTALLED:
use_test_ot_exporter()
cls.ot_exporter = get_test_ot_exporter()
def teardown_method(self):
if HAS_OPENTELEMETRY_INSTALLED:
self.ot_exporter.clear()
def assertNoSpans(self):
if HAS_OPENTELEMETRY_INSTALLED:
span_list = self.ot_exporter.get_finished_spans()
self.assertEqual(len(span_list), 0)
def assertSpanAttributes(
self, name, status=StatusCode.OK, attributes=None, span=None
):
if HAS_OPENTELEMETRY_INSTALLED:
if not span:
span_list = self.ot_exporter.get_finished_spans()
self.assertEqual(len(span_list), 1)
span = span_list[0]
self.assertEqual(span.name, name)
self.assertEqual(span.status.status_code, status)
self.assertEqual(dict(span.attributes), attributes)