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')
