Coverage for webapp/store_api.py: 100%

1import talisker 

2from canonicalwebteam.store_api.devicegw import DeviceGW 

3from canonicalwebteam.store_api.publishergw import PublisherGW 

4from webapp.observability.utils import trace_function 

5 

6 

7def decorate_all_methods(decorator, cls): 

8 class WrappedClass(cls): 

9 pass 

10 

11 for attr_name in dir(cls): 

12 # ignore private and special methods like __init__ etc. 

13 if attr_name.startswith("__"): 

14 continue 

15 attr = getattr(cls, attr_name) 

16 # only decorate callable attributes (these are mostly API calls) 

17 if callable(attr): 

18 setattr(WrappedClass, attr_name, decorator(attr)) 

19 

20 return WrappedClass 

21 

22 

23TracedPubilsherGW = decorate_all_methods(trace_function, PublisherGW) 

24TracedDeviceGW = decorate_all_methods(trace_function, DeviceGW) 

25 

26request_session = talisker.requests.get_session() 

27publisher_gateway = TracedPubilsherGW("charm", request_session) 

28device_gateway = TracedDeviceGW("charm", request_session)