Coverage for osm_policy_module/common/message_bus_client.py: 89%

1# -*- coding: utf-8 -*- 

2 

3# Copyright 2018 Whitestack, LLC 

4# ************************************************************* 

5 

6# This file is part of OSM Monitoring module 

7# All Rights Reserved to Whitestack, LLC 

8 

9# Licensed under the Apache License, Version 2.0 (the "License"); you may 

10# not use this file except in compliance with the License. You may obtain 

11# a copy of the License at 

12 

13# http://www.apache.org/licenses/LICENSE-2.0 

14 

15# Unless required by applicable law or agreed to in writing, software 

16# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT 

17# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 

18# License for the specific language governing permissions and limitations 

19# under the License. 

20 

21# For those usages not covered by the Apache License, Version 2.0 please 

22# contact: bdiaz@whitestack.com or glavado@whitestack.com 

23## 

24from typing import List, Callable 

25 

26from osm_common import msgkafka, msglocal 

27 

28from osm_policy_module.core.config import Config 

29 

30 

31class MessageBusClient: 

32 def __init__(self, config: Config): 

33 if config.get("message", "driver") == "local": 

34 self.msg_bus = msglocal.MsgLocal() 

35 elif config.get("message", "driver") == "kafka": 

36 self.msg_bus = msgkafka.MsgKafka() 

37 else: 

38 raise Exception( 

39 "Unknown message bug driver {}".format(config.get("section", "driver")) 

40 ) 

41 self.msg_bus.connect(config.get("message")) 

42 

43 async def aioread(self, topics: List[str], callback: Callable = None, **kwargs): 

44 """ 

45 Retrieves messages continuously from bus and executes callback for each message consumed. 

46 :param topics: List of message bus topics to consume from. 

47 :param callback: Async callback function to be called for each message received. 

48 :param kwargs: Keyword arguments to be passed to callback function. 

49 :return: None 

50 """ 

51 await self.msg_bus.aioread(topics, aiocallback=callback, **kwargs) 

52 

53 async def aiowrite(self, topic: str, key: str, msg: dict): 

54 """ 

55 Writes message to bus. 

56 :param topic: Topic to write to. 

57 :param key: Key to write to. 

58 :param msg: Dictionary containing message to be written. 

59 :return: None 

60 """ 

61 await self.msg_bus.aiowrite(topic, key, msg) 

62 

63 async def aioread_once(self, topic: str): 

64 """ 

65 Retrieves last message from bus. 

66 :param topic: topic to retrieve message from. 

67 :return: tuple(topic, key, message) 

68 """ 

69 result = await self.msg_bus.aioread(topic) 

70 return result