Skip to content

Commit

Permalink
fix linting
Browse files Browse the repository at this point in the history
  • Loading branch information
melvinkcx committed Aug 15, 2022
1 parent d0f99a9 commit 9032073
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 17 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,26 +12,25 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from inspect import ismodule, getmembers, isclass
from inspect import getmembers, isclass, ismodule
from typing import Collection

import fastapi_events
import wrapt
from fastapi_events.handlers.base import BaseEventHandler
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.instrumentation.utils import unwrap
from opentelemetry.trace import SpanKind

from opentelemetry import trace
from opentelemetry.instrumentation.fastapi_events.package import _instruments
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.instrumentation.utils import unwrap
from opentelemetry.trace import SpanKind


async def _handle_wrapper(wrapped, instance, args, kwargs):
tracer = trace.get_tracer(__name__)
event = args[0] if args else kwargs.get("event")
with tracer.start_as_current_span(
f"handling event {event[0]}",
kind=SpanKind.CONSUMER
f"handling event {event[0]}", kind=SpanKind.CONSUMER
) as span:
return await wrapped(event)

Expand All @@ -40,8 +39,7 @@ async def _handle_many_wrapper(wrapped, instance, args, kwargs):
tracer = trace.get_tracer(__name__)
events = args[0] if args else kwargs.get("events")
with tracer.start_as_current_span(
f"handling multiple events",
kind=SpanKind.CONSUMER
f"handling multiple events", kind=SpanKind.CONSUMER
) as span:
return await wrapped(events)

Expand All @@ -68,12 +66,16 @@ def _instrument(self, **kwargs):
for _, class_ in getmembers(module, isclass):
if issubclass(class_, BaseEventHandler):
self._instrumented_classes.append(class_)
wrapt.wrap_function_wrapper(class_, "handle", _handle_wrapper)
wrapt.wrap_function_wrapper(class_, "handle_many", _handle_many_wrapper)

wrapt.wrap_function_wrapper(fastapi_events.dispatcher,
"dispatch",
dispatch_wrapper)
wrapt.wrap_function_wrapper(
class_, "handle", _handle_wrapper
)
wrapt.wrap_function_wrapper(
class_, "handle_many", _handle_many_wrapper
)

wrapt.wrap_function_wrapper(
fastapi_events.dispatcher, "dispatch", dispatch_wrapper
)

def _uninstrument(self, **kwargs):
for class_ in self._instrumented_classes:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,12 @@
from fastapi_events.dispatcher import dispatch
from fastapi_events.handlers.local import LocalHandler
from fastapi_events.middleware import EventHandlerASGIMiddleware
from opentelemetry.test.test_base import TestBase
from starlette.testclient import TestClient

from opentelemetry.instrumentation.fastapi_events import FastAPIEventsInstrumentor
from opentelemetry.instrumentation.fastapi_events import (
FastAPIEventsInstrumentor,
)
from opentelemetry.test.test_base import TestBase


class TestFastAPIEventsInstrumentor(TestBase):
Expand All @@ -32,7 +34,9 @@ def _create_app(self):
app = FastAPI()
local_handler = LocalHandler()

app.add_middleware(EventHandlerASGIMiddleware, handlers=[local_handler])
app.add_middleware(
EventHandlerASGIMiddleware, handlers=[local_handler]
)

FastAPIEventsInstrumentor().instrument()

Expand Down

0 comments on commit 9032073

Please sign in to comment.