class Collector(object): def __init__(self, dst_path): self.downloader = AuditsDownloader(dst_path) self.is_downloading = False self.etcd_client = Client(address=ETCD_SERVICE_HOST, port='2379') with self.blocking_channel() as channel: # Make sure queue is there channel.queue_declare(queue=QUEUE_NAME) logging.info('queue {} declared'.format(QUEUE_NAME)) @contextmanager def blocking_channel(self): connection = pika.BlockingConnection( pika.ConnectionParameters(RABBITMQ_SERVICE_HOST)) channel = connection.channel() yield channel channel.close() connection.close() def add_to_queue(self, filepath): with self.blocking_channel() as channel: channel.basic_publish(exchange='', routing_key=QUEUE_NAME, body=filepath) logging.info('Placed in queue file {}'.format(filepath)) def set_url(self): try: url = self.etcd_client.get('/data/collector/url').value if url.lower() not in ['none', 'null']: self.downloader.url = url except KeyError: logging.info('No hdfs address set yet') @contextmanager def downloading(self): self.is_downloading = True yield self.is_downloading = False @property def delay(self): try: return int(self.etcd_client.get('/data/collector/delay').value) except KeyError: return 60 def run(self): self.set_url() with self.downloading(): self.downloader.download(callback=self.add_to_queue) def run_periodically(self): self.run() scheduler = sched.scheduler(time.time, time.sleep) while True: if not self.is_downloading: # Add to queue scheduler.enter(self.delay, 1, self.run, []) scheduler.run()
def test_get_raises_an_error_if_directory(): requester = Mock() requester.get = Mock(side_effect=get_response) client = Client(requester=requester) with pytest.raises(errors.NotAFile): client.get('/k')