1212# See the License for the specific language governing permissions and
1313# limitations under the License.
1414
15+ import asyncio
16+
17+ from tortoise import Tortoise , fields , models
18+
19+ from opentelemetry import trace
1520from opentelemetry .instrumentation .tortoiseorm import TortoiseORMInstrumentor
21+ from opentelemetry .instrumentation .utils import suppress_instrumentation
22+ from opentelemetry .semconv ._incubating .attributes .db_attributes import (
23+ DB_NAME ,
24+ DB_STATEMENT ,
25+ DB_SYSTEM ,
26+ )
1627from opentelemetry .test .test_base import TestBase
1728
18- # pylint: disable=too-many-public-methods
29+
30+ class MockModel (models .Model ):
31+ id = fields .IntField (pk = True )
32+ name = fields .TextField ()
33+
34+ def __str__ (self ):
35+ return self .name
1936
2037
2138class TestTortoiseORMInstrumentor (TestBase ):
@@ -26,9 +43,94 @@ def setUp(self):
2643 def tearDown (self ):
2744 super ().tearDown ()
2845 TortoiseORMInstrumentor ().uninstrument ()
46+ self ._async_call (Tortoise ._drop_databases ())
47+
48+ # pylint: disable-next=no-self-use
49+ def _async_call (self , coro ):
50+ return asyncio .run (coro )
51+
52+ # pylint: disable-next=no-self-use
53+ async def _init_tortoise (self ):
54+ await Tortoise .init (
55+ db_url = "sqlite://:memory:" ,
56+ modules = {"models" : [__name__ ]},
57+ )
58+ await Tortoise .generate_schemas ()
59+
60+ def test_trace_integration (self ):
61+ async def run ():
62+ await self ._init_tortoise ()
63+ await MockModel .create (name = "Test 1" )
64+ await MockModel .filter (name = "Test 1" ).first ()
65+
66+ self ._async_call (run ())
67+ spans = self .memory_exporter .get_finished_spans ()
68+
69+ crud_spans = [s for s in spans if s .name in ("INSERT" , "SELECT" )]
70+ self .assertGreaterEqual (len (crud_spans ), 2 )
71+
72+ insert_span = next (s for s in crud_spans if s .name == "INSERT" )
73+ self .assertEqual (insert_span .kind , trace .SpanKind .CLIENT )
74+ self .assertEqual (insert_span .attributes [DB_SYSTEM ], "sqlite" )
75+ self .assertEqual (insert_span .attributes [DB_NAME ], ":memory:" )
76+ self .assertIn ("INSERT" , insert_span .attributes [DB_STATEMENT ])
77+
78+ select_span = next (s for s in crud_spans if s .name == "SELECT" )
79+ self .assertEqual (select_span .kind , trace .SpanKind .CLIENT )
80+ self .assertEqual (select_span .attributes [DB_SYSTEM ], "sqlite" )
81+ self .assertEqual (select_span .attributes [DB_NAME ], ":memory:" )
82+ self .assertIn ("SELECT" , select_span .attributes [DB_STATEMENT ])
83+
84+ def test_capture_parameters (self ):
85+ TortoiseORMInstrumentor ().uninstrument ()
86+ TortoiseORMInstrumentor ().instrument (capture_parameters = True )
87+
88+ async def run ():
89+ await self ._init_tortoise ()
90+ await MockModel .create (name = "Test Parameterized" )
91+
92+ self ._async_call (run ())
93+ spans = self .memory_exporter .get_finished_spans ()
94+ insert_span = next (s for s in spans if s .name == "INSERT" )
95+ self .assertIn ("db.statement.parameters" , insert_span .attributes )
96+ self .assertIn (
97+ "Test Parameterized" ,
98+ insert_span .attributes ["db.statement.parameters" ],
99+ )
100+
101+ def test_uninstrument (self ):
102+ TortoiseORMInstrumentor ().uninstrument ()
103+
104+ async def run ():
105+ await self ._init_tortoise ()
106+ await MockModel .create (name = "Test Uninstrumented" )
107+
108+ self ._async_call (run ())
109+ spans = self .memory_exporter .get_finished_spans ()
110+ crud_spans = [s for s in spans if s .name in ("INSERT" , "SELECT" )]
111+ self .assertEqual (len (crud_spans ), 0 )
112+
113+ def test_suppress_instrumentation (self ):
114+ async def run ():
115+ await self ._init_tortoise ()
116+ with suppress_instrumentation ():
117+ await MockModel .create (name = "Test Suppressed" )
118+
119+ self ._async_call (run ())
120+ spans = self .memory_exporter .get_finished_spans ()
121+ crud_spans = [s for s in spans if s .name in ("INSERT" , "SELECT" )]
122+ self .assertEqual (len (crud_spans ), 0 )
123+
124+ def test_no_op_tracer_provider (self ):
125+ TortoiseORMInstrumentor ().uninstrument ()
126+ TortoiseORMInstrumentor ().instrument (
127+ tracer_provider = trace .NoOpTracerProvider ()
128+ )
129+
130+ async def run ():
131+ await self ._init_tortoise ()
132+ await MockModel .create (name = "Test NoOp" )
29133
30- def test_tortoise (self ):
31- # FIXME This instrumentation has no tests at all and should have some
32- # tests. This is being added just to make pytest not fail because no
33- # tests are being collected for tortoise at the moment.
34- pass
134+ self ._async_call (run ())
135+ spans = self .memory_exporter .get_finished_spans ()
136+ self .assertEqual (len (spans ), 0 )
0 commit comments