Increase timeout for member stop
[integration/test.git] / tools / wstools / wsreceiver.py
1 """WebSocket data receiver.
2
3 The tool receives and logs data from specified URI"""
4
5 # Copyright (c) 2016 Cisco Systems, Inc. and others.  All rights reserved.
6 #
7 # This program and the accompanying materials are made available under the
8 # terms of the Eclipse Public License v1.0 which accompanies this distribution,
9 # and is available at http://www.eclipse.org/legal/epl-v10.html
10
11 import argparse
12 import logging
13 from websocket import create_connection
14
15 __author__ = "Radovan Sajben"
16 __copyright__ = "Copyright(c) 2016, Cisco Systems, Inc."
17 __license__ = "Eclipse Public License v1.0"
18 __email__ = "rsajben@cisco.com"
19
20
21 def parse_arguments():
22     """Use argparse to get arguments,
23
24     Returns:
25         :return: args object
26     """
27     parser = argparse.ArgumentParser()
28     parser.add_argument("--uri", default="ws://127.0.0.1:8185/", help="URI to connect")
29     parser.add_argument("--count", type=int, default=1, help="Number of messages to receive")
30     parser.add_argument("--logfile", default="wsreceiver.log", help="Log file name")
31     parser.add_argument("--debug", dest="loglevel", action="store_const",
32                         const=logging.DEBUG, default=logging.INFO, help="Log level")
33     args = parser.parse_args()
34     return args
35
36
37 class WSReceiver(object):
38     """Class which receives web socket messages."""
39
40     def __init__(self, uri):
41         """Initialise and connect to URI.
42
43         Arguments:
44             :uri: uri to connect to
45         Returns:
46             None
47         """
48         self.uri = uri
49         logger.info("Connecting to: %s", self.uri)
50         self.ws = create_connection(self.uri)
51
52     def close(self):
53         """Close the connection.
54
55         Arguments:
56             None
57         Returns:
58             None
59         """
60         logger.info("Disconnecting from: %s", self.uri)
61         self.ws.close()
62
63     def receive(self):
64         """Receive a message.
65
66         Arguments:
67             None
68         Returns:
69             :return: received data
70         """
71         data = self.ws.recv()
72         logger.info("Data received:\n%s", data)
73         return data
74
75
76 if __name__ == "__main__":
77     args = parse_arguments()
78     logger = logging.getLogger("logger")
79     log_formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
80     console_handler = logging.StreamHandler()
81     file_handler = logging.FileHandler(args.logfile, mode="w")
82     console_handler.setFormatter(log_formatter)
83     file_handler.setFormatter(log_formatter)
84     logger.addHandler(console_handler)
85     logger.addHandler(file_handler)
86     logger.setLevel(args.loglevel)
87     receiver = WSReceiver(args.uri)
88     remains = args.count
89     logger.info("Expected %d message(s)", remains)
90     while remains:
91         logger.info("Waiting for a message ...")
92         data = receiver.receive()
93         remains -= 1
94         logger.info("Remaining messages to receive: %d", remains)
95     logger.info("Finished ...")
96     receiver.close()