|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
14 | 14 |
|
| 15 | +import asyncio |
15 | 16 | import types |
16 | 17 | from unittest import IsolatedAsyncioTestCase, mock |
17 | 18 |
|
@@ -50,10 +51,15 @@ def __init__(self, *args, **kwargs): |
50 | 51 | pass |
51 | 52 |
|
52 | 53 | # pylint: disable=unused-argument, no-self-use |
53 | | - async def execute(self, query, params=None, throw_exception=False): |
| 54 | + async def execute( |
| 55 | + self, query, params=None, throw_exception=False, delay=0.0 |
| 56 | + ): |
54 | 57 | if throw_exception: |
55 | 58 | raise psycopg.Error("Test Exception") |
56 | 59 |
|
| 60 | + if delay: |
| 61 | + await asyncio.sleep(delay) |
| 62 | + |
57 | 63 | # pylint: disable=unused-argument, no-self-use |
58 | 64 | async def executemany(self, query, params=None, throw_exception=False): |
59 | 65 | if throw_exception: |
@@ -492,3 +498,27 @@ async def test_not_recording_async(self): |
492 | 498 | self.assertFalse(mock_span.set_status.called) |
493 | 499 |
|
494 | 500 | PsycopgInstrumentor().uninstrument() |
| 501 | + |
| 502 | + async def test_tracing_is_async(self): |
| 503 | + PsycopgInstrumentor().instrument() |
| 504 | + |
| 505 | + # before this async fix cursor.execute would take 14000 ns, delaying for |
| 506 | + # 100,000ns |
| 507 | + delay = 0.0001 |
| 508 | + |
| 509 | + async def test_async_connection(): |
| 510 | + acnx = await psycopg.AsyncConnection.connect("test") |
| 511 | + async with acnx as cnx: |
| 512 | + async with cnx.cursor() as cursor: |
| 513 | + await cursor.execute("SELECT * FROM test", delay=delay) |
| 514 | + |
| 515 | + await test_async_connection() |
| 516 | + spans_list = self.memory_exporter.get_finished_spans() |
| 517 | + self.assertEqual(len(spans_list), 1) |
| 518 | + span = spans_list[0] |
| 519 | + |
| 520 | + # duration is nanoseconds |
| 521 | + duration = span.end_time - span.start_time |
| 522 | + self.assertGreater(duration, delay * 1e9) |
| 523 | + |
| 524 | + PsycopgInstrumentor().uninstrument() |
0 commit comments