Skip to content

Commit

Permalink
datadog: fix test_memory_trace_client_test2.py
Browse files Browse the repository at this point in the history
  • Loading branch information
naisanzaa committed Apr 13, 2024
1 parent 4400120 commit 02fcdbe
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 32 deletions.
32 changes: 22 additions & 10 deletions automon/integrations/openTelemetryWrapper/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,18 +16,18 @@ class OpenTelemetryClient(object):
def __init__(self):
self.config = OpenTelemetryConfig()

async def add_event(self, name: str, attributes: types.Attributes = None, **kwargs):
def add_event(self, name: str, attributes: types.Attributes = None, **kwargs):
logger.debug(dict(name=name, attributes=attributes, kwargs=kwargs))
span = await self.current_span()
span = self.current_span()
return span.add_event(name=name, attributes=attributes, **kwargs)

async def clear(self):
logger.debug('clear')
return await self.config.clear()

async def current_span(self):
def current_span(self):
logger.debug('current_span')
return await self.config.current_span()
return self.config.current_span()

async def is_ready(self):
if await self.config.is_ready():
Expand All @@ -47,7 +47,7 @@ async def record_exception(self, exception: Exception):
span.set_status(Status(StatusCode.ERROR))
return span.record_exception(exception=exception)

async def start_as_current_span(
def start_as_current_span(
self, name: str,
attributes: types.Attributes = None,
**kwargs
Expand All @@ -72,13 +72,25 @@ async def start_producer(self):

async def test(self):
with await self.start_as_current_span(name='rootSpan') as trace_root:
await self.add_event('AAAAAAAA')
self.config.opentelemetry.trace.get_current_span().add_event('AAAAAAAA')

with self.config.tracer.start_as_current_span(name='childSpan') as trace_child:
self.config.opentelemetry.trace.get_current_span().add_event('AAAAAAAA')
self.config.opentelemetry.trace.get_current_span().add_event('BBBBBBBB')

self.config.opentelemetry.trace.get_current_span().add_event('BBBBBBBB')

return True

async def test2(self):
with self.config.tracer.start_as_current_span(name='rootSpan') as trace_root:
self.config.opentelemetry.trace.get_current_span().add_event('AAAAAAAA')

with await self.start_as_current_span(name='childSpan') as trace_child:
await self.add_event('AAAAAAAA')
await self.add_event('BBBBBBBB')
with self.config.tracer.start_as_current_span(name='childSpan') as trace_child:
self.config.opentelemetry.trace.get_current_span().add_event('AAAAAAAA')
self.config.opentelemetry.trace.get_current_span().add_event('BBBBBBBB')

await self.add_event('BBBBBBBB')
self.config.opentelemetry.trace.get_current_span().add_event('BBBBBBBB')

return True

Expand Down
3 changes: 2 additions & 1 deletion automon/integrations/openTelemetryWrapper/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

class OpenTelemetryConfig(object):
def __init__(self):
self.opentelemetry = opentelemetry
self.provider = TracerProvider()
self.memory_processor = InMemorySpanExporter()
self.processor = SimpleSpanProcessor(self.memory_processor)
Expand All @@ -29,7 +30,7 @@ async def clear(self):
logger.debug('clear')
return self.memory_processor.clear()

async def current_span(self):
def current_span(self):
logger.debug('get_current_span')
return opentelemetry.trace.get_current_span()

Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
import unittest
import asyncio

from automon.integrations.openTelemetryWrapper import OpenTelemetryClient


class MyTestCase(unittest.TestCase):
client = OpenTelemetryClient()

def test_something(self):
self.assertTrue(asyncio.run(
self.client.test()))

pass


if __name__ == '__main__':
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
import unittest
import asyncio

from automon.integrations.openTelemetryWrapper import OpenTelemetryClient


class MyTestCase(unittest.TestCase):
client = OpenTelemetryClient()

def test_something(self):
self.assertTrue(asyncio.run(
self.client.test2()))

spans = asyncio.run(
self.client.get_finished_spans())

pass


if __name__ == '__main__':
unittest.main()

0 comments on commit 02fcdbe

Please sign in to comment.