Documents up-to-date
[doctor.git] / doctor_tests / consumer / sample.py
1 ##############################################################################
2 # Copyright (c) 2017 ZTE Corporation and others.
3 #
4 # All rights reserved. This program and the accompanying materials
5 # are made available under the terms of the Apache License, Version 2.0
6 # which accompanies this distribution, and is available at
7 # http://www.apache.org/licenses/LICENSE-2.0
8 ##############################################################################
9 from flask import Flask
10 from flask import request
11 import json
12 import time
13 from threading import Thread
14 import requests
15
16 from doctor_tests.consumer.base import BaseConsumer
17
18
19 class SampleConsumer(BaseConsumer):
20
21     def __init__(self, conf, log):
22         super(SampleConsumer, self).__init__(conf, log)
23         self.app = None
24
25     def start(self):
26         self.log.info('sample consumer start......')
27         self.app = ConsumerApp(self.conf.consumer.port, self, self.log)
28         self.app.start()
29
30     def stop(self):
31         self.log.info('sample consumer stop......')
32         if not self.app:
33             return
34         headers = {
35             'Content-Type': 'application/json',
36             'Accept': 'application/json',
37         }
38         url = 'http://%s:%d/shutdown'\
39               % (self.conf.consumer.ip,
40                  self.conf.consumer.port)
41         requests.post(url, data='', headers=headers)
42
43
44 class ConsumerApp(Thread):
45
46     def __init__(self, port, consumer, log):
47         Thread.__init__(self)
48         self.port = port
49         self.consumer = consumer
50         self.log = log
51
52     def run(self):
53         app = Flask('consumer')
54
55         @app.route('/failure', methods=['POST'])
56         def event_posted():
57             notified_time = time.time()
58             self.log.info('doctor consumer notified at %s' % notified_time)
59             self.consumer.notified_time = notified_time
60             data = json.loads(request.data.decode('utf8'))
61             self.log.info('sample consumer received data = %s' % data)
62             return 'OK'
63
64         @app.route('/shutdown', methods=['POST'])
65         def shutdown():
66             self.log.info('shutdown consumer app server at %s' % time.time())
67             func = request.environ.get('werkzeug.server.shutdown')
68             if func is None:
69                 raise RuntimeError('Not running with the Werkzeug Server')
70             func()
71             return 'consumer app shutting down...'
72
73         app.run(host="0.0.0.0", port=self.port)