diff --git a/swagger_zipkin/zipkin_decorator.py b/swagger_zipkin/zipkin_decorator.py index d6146bf..f147bbd 100644 --- a/swagger_zipkin/zipkin_decorator.py +++ b/swagger_zipkin/zipkin_decorator.py @@ -20,6 +20,9 @@ def __init__(self, resource): def __getattr__(self, name): return decorate_client(self.resource, self.with_headers, name) + def __call__(self, *args, **kwargs): + return self.resource(*args, **kwargs) + def with_headers(self, call_name, *args, **kwargs): kwargs.setdefault('_request_options', {}) headers = kwargs['_request_options'].setdefault('headers', {}) diff --git a/tests/zipkin_decorator_test.py b/tests/zipkin_decorator_test.py index cf3b815..cf2137b 100644 --- a/tests/zipkin_decorator_test.py +++ b/tests/zipkin_decorator_test.py @@ -29,6 +29,22 @@ def create_request_options(trace_id, span_id, parent_span_id): } +def create_client(): + class V1Operation(): + def foo(): + return 'foo' + + class Client(object): + def __init__(self): + self.v1 = V1Operation() + + def func(self): + return 'some func' + + client = Client() + return client + + def test_client_request_option_decorator(): trace_id = 'trace_id' span_id = 'span_id' @@ -60,15 +76,13 @@ def test_client_request_option_decorator(): def test_client_dir(): - class V1Operation(): - def foo(): - return 'foo' - - class Client(object): - def __init__(self): - self.v1 = V1Operation() - - client = Client() + client = create_client() wrapped_client = ZipkinClientDecorator(client) assert 'foo' in dir(wrapped_client.v1) + + +def test_callable(): + client = create_client() + wrapped_client = ZipkinClientDecorator(client) + assert 'some func' == wrapped_client.func()