def testPubSub(self):
sub_data = {'secret': 'SECRET'}
id = PubSub.subscribe ('EVENT', 'not_an_url', sub_data)
self.assertTrue (PubSub.unsubscribe (id))
self.assertFalse (PubSub.unsubscribe (id))
pub_data = {'message': 123}
self.assertEqual(0, PubSub.publish ('EVENT', pub_data))
PubSub.subscribe ('EVENT', 'not_an_url', sub_data)
self.assertEqual(1, PubSub.publish ('EVENT', pub_data))
response = self.executeTask() # /pubsub/notify
self.assertEqual(response.status_int, 200)
response.mustcontain ('unknown url type')
url = "/pubsub/test"
PubSub.subscribe ('EVENT2', url, sub_data)
self.assertEqual(1, PubSub.publish ('EVENT2', pub_data))
response = self.executeTask() # /pubsub/notify
self.assertEqual (response.json['status'], 'OK')
response = self.executeTask() # /pubsub/test
self.assertEqual (response.json['pub_data']["message"], 123)
url = "/pubsub/task"
sub_data = {'secret': 'SECRET', 'channel': 'CHANNEL', 'taskname': 'NAME'}
PubSub.subscribe ('EVENT3', url, sub_data)
self.assertEqual(1, PubSub.publish ('EVENT3', pub_data))
response = self.executeTask() # /pubsub/notify
self.assertEqual (response.json['status'], 'OK')
response = self.executeTask() # /pubsub/task
self.assertEqual (response.json['status'], 'OK')
queue = TaskQueue()
lease = queue.lease (channel='CHANNEL')
self.assertIsNotNone(lease)
self.assertEqual(lease['id'], response.json['id'])
def execute(tag, queue, server, qurl):
tq = TaskQueue(queue_name=queue,
queue_server=server,
n_threads=0,
qurl=qurl)
print("Pulling from {}://{}".format(server, queue))
tries = 0
with tq:
while LOOP:
task = 'unknown'
try:
task = tq.lease(tag=tag, seconds=int(LEASE_SECONDS))
tries += 1
print(task)
task.execute()
tq.delete(task)
logger.log('INFO', task, "succesfully executed")
tries = 0
except TaskQueue.QueueEmpty:
time.sleep(random_exponential_window_backoff(tries))
continue
except EmptyVolumeException:
logger.log('WARNING', task, "raised an EmptyVolumeException")
tq.delete(task)
except Exception as e:
logger.log('ERROR', task,
"raised {}\n {}".format(e, traceback.format_exc()))
raise #this will restart the container in kubernetes
def mv(src, dest):
"""
Moves the contents of a queue to another
service or location. Do not run this
process while a queue is being worked.
Moving an sqs queue to a file queue
may result in duplicated tasks.
"""
src = normalize_path(src)
dest = normalize_path(dest)
tqd = TaskQueue(dest, progress=False)
tqs = TaskQueue(src, progress=False)
total = tqs.enqueued
with tqdm(total=total, desc="Moving") as pbar:
while True:
try:
tasks = tqs.lease(num_tasks=10, seconds=10)
except QueueEmptyError:
break
tqd.insert(tasks)
tqs.delete(tasks)
pbar.update(len(tasks))
def lease(self, request):
"""Lease a task in the queue for a specified period of time
Args:
request: a LeaseTaskRequest
lease_seconds is the number of seconds to wait for the task to be completed. If the task is not deleted from the queue
in that interval, then it will become available again for lease
Returns:
A LeaseTaskResponse. status == 'OK' on success. If there are no tasks available in the requested channel, then
status == 'NOT FOUND'. Any other value for status indicates failure.
the task_id is a unique identifier assigned to the task request when it is added. It is needed in order to delete the task
the task attribute in the response is a copy of the AddTaskRequest message that was supplied in the call to taskqueue.add
"""
response = LeaseTaskResponse(status='NOT FOUND', task_id='')
try:
EndpointsHelper.authenticate()
task = TaskQueue.lease(request.channel, request.lease_seconds)
if task:
response.status='OK'
response.task = AddTaskRequest(channel=task['channel'], name=task['name'], content=task['content'])
response.task_id = task['id']
response.retry_count = task['retry_count']
except Exception, err:
response.status=str(err)
def testretry(self):
queue = TaskQueue()
id = queue.add (channel='RERTY_TEST', name='N', content=dict(xxx=123))
self.assertIsNotNone(id)
lease = queue.lease (channel='RERTY_TEST', lease_seconds=1)
self.assertEqual(id, lease['id'])
self.assertEqual(lease['retry_count'], 1)
time.sleep(1)
lease = queue.lease (channel='RERTY_TEST', lease_seconds=1)
self.assertEqual(id, lease['id'])
self.assertEqual(lease['retry_count'], 2)
time.sleep(1)
lease = queue.lease (channel='RERTY_TEST', lease_seconds=1)
self.assertEqual(id, lease['id'])
self.assertEqual(lease['retry_count'], 3)
time.sleep(1)
lease = queue.lease (channel='RERTY_TEST', lease_seconds=1)
self.assertIsNone(lease)
def test_blackout_tasks():
delete_layer()
cf, _ = create_layer(size=(128,64,64,1), offset=(0,0,0), layer_type="image")
cv = CloudVolume(cf.cloudpath)
tq = TaskQueue("fq:///tmp/removeme/blackout/")
tq.insert(
partial(BlackoutTask,
cloudpath=cf.cloudpath,
mip=0,
offset=(0,0,0),
shape=(128, 64, 64),
value=11,
non_aligned_writes=False
)
)
tq.lease().execute()
img = cv[:,:,:]
assert np.all(img == 11)
BlackoutTask(
cloudpath=cf.cloudpath,
mip=0,
offset=(0,0,0),
shape=(37, 64, 64),
value=23,
non_aligned_writes=True
)
img = cv[:37,:,:]
assert np.all(img == 23)
img = cv[:]
items, counts = np.unique(img, return_counts=True)
counts = {
items[0]: counts[0],
items[1]: counts[1]
}
twenty_threes = 37 * 64 * 64
assert counts[23] == twenty_threes
assert counts[11] == (128 * 64 * 64) - twenty_threes
def test_get(sqs, protocol):
path = getpath(protocol)
tq = TaskQueue(path, n_threads=0)
n_inserts = 5
tq.purge()
tq.insert((PrintTask() for _ in range(n_inserts)))
for i in range(n_inserts):
t = tq.lease()
tq.delete(t)
def test_lease(sqs):
path = getpath("sqs")
tq = TaskQueue(path, n_threads=0)
n_inserts = 20
tq.purge()
tq.insert(( PrintTask(str(x)) for x in range(n_inserts) ))
tasks = tq.lease(num_tasks=10, wait_sec=0)
assert len(tasks) == 10
tq.delete(tasks)
tasks = tq.lease(num_tasks=10, wait_sec=0)
assert len(tasks) == 10
tq.delete(tasks)
try:
tasks = tq.lease(num_tasks=10, wait_sec=0)
assert False
except QueueEmptyError:
pass
def execute(tag, queue, seconds, loop):
tq = TaskQueue(queue)
print("Pulling from {}".format(tq.qualified_path))
seconds = int(seconds)
if loop:
tq.poll(lease_seconds=seconds, verbose=True)
else:
task = tq.lease(seconds=seconds)
task.execute()
def test_is_empty():
tq = TaskQueue(FILE_QURL)
tq.purge()
assert tq.is_empty() == True
tq.insert(PrintTask("hello"))
assert tq.is_empty() == False
task = tq.lease()
tq.delete(task)
assert tq.is_empty() == True
def execute(tag, queue, server, qurl, loop):
tq = TaskQueue(queue_name=queue,
queue_server=server,
n_threads=0,
qurl=qurl)
print("Pulling from {}://{}".format(server, queue))
sec = int(LEASE_SECONDS)
if loop:
tq.poll(lease_seconds=sec, verbose=True)
else:
task = tq.lease(seconds=sec)
task.execute()
def testGeoFeed(self):
item = {
'topic': 'T',
'key': 'K',
'latitude': 39,
'longitude': -79
}
id = GeoFeed.publish(**item)
self.assertEqual(id, item['key'])
time.sleep(2)
item['key'] = 'L'
id = GeoFeed.publish(**item)
self.assertEqual(id, item['key'])
last_published = '9999'
for doc in GeoFeed.list(item['topic']):
self.assertLess(doc['published'], last_published)
last_published = doc['published']
for doc in GeoFeed.list('NOT_FOUND'):
self.assertFalse('Should never get here')
doc = GeoFeed.get(item['topic'], item['key'])
self.assertIsNotNone(doc)
#set up PubSub subscription so that a task is created when an item is published to the feed
sub_url = "/pubsub/task"
event=GeoFeed._indexname(item['topic'])
channel='ProcessNew%s' % item['topic']
sub_data = {'secret': 'SECRET', 'channel': channel, 'pubname': 'key'}
PubSub.subscribe (event, sub_url, sub_data)
# now publish a new item to the feed. This should trigger creation of a new task in the queue
id = GeoFeed.publish(**item)
# need to manually process the task queue because we're in test mode
response = self.executeTask() # /pubsub/notify
self.assertEqual (response.json['status'], 'OK')
response = self.executeTask() # /pubsub/task
self.assertEqual (response.json['status'], 'OK')
# Now make sure there is a task in the queue
queue = TaskQueue()
lease = queue.lease (channel=channel)
self.assertIsNotNone(lease)
self.assertEqual(lease['id'], response.json['id'])
self.assertEqual(lease['content']['pub_data']['key'], item['key'])
def execute_helper(queue, aws_region, lease_sec, tally, min_sec):
tq = TaskQueue(normalize_path(queue), region_name=aws_region)
def stop_after_elapsed_time(elapsed_time):
if min_sec < 0:
return False
return min_sec < elapsed_time
if min_sec != 0:
tq.poll(
lease_seconds=lease_sec,
verbose=True,
tally=tally,
stop_fn=stop_after_elapsed_time,
)
else:
task = tq.lease(seconds=lease_sec)
task.execute()
def test_queue_transfer(sqs):
tqsqs = TaskQueue(getpath("sqs"))
tqsqs.purge()
tqfq = TaskQueue(getpath("fq"))
tqfq.purge()
assert tqsqs.enqueued == 0
tqfq.insert((PrintTask() for _ in range(10)))
tqsqs.insert(tqfq)
assert tqsqs.enqueued == 10
task = tqsqs.lease()
assert isinstance(task, PrintTask)
try:
tqfq.insert(tqsqs)
assert False
except taskqueue.UnsupportedProtocolError:
pass
def execute(queue, region_name, endpoint_url, seconds, tally, loop):
tq = TaskQueue(queue, region_name=region_name, endpoint_url=endpoint_url)
print("Pulling from {}".format(tq.qualified_path))
seconds = int(seconds)
def stop_after_elapsed_time(elapsed_time):
if loop < 0:
return False
return loop < elapsed_time
if loop != 0:
tq.poll(
lease_seconds=seconds,
verbose=True,
tally=tally,
stop_fn=stop_after_elapsed_time,
)
else:
task = tq.lease(seconds=seconds)
task.execute()
def testall(self):
queue = TaskQueue()
id = queue.add (channel='C', name='N', content=dict(xxx=123))
self.assertIsNotNone(id)
lease = queue.lease (channel='C')
self.assertIsNotNone(lease)
self.assertEqual(lease['id'], id)
self.assertEqual(lease['name'], 'N')
self.assertEqual(lease['channel'], 'C')
lease = queue.lease (channel='C')
self.assertIsNone(lease)
lease = queue.lease (channel='NOT FOUND')
self.assertIsNone(lease)
queue.add (channel='C', name='N1', content=dict(xxx=123))
queue.add (channel='C', name='N2', content=dict(xxx=123))
lease = queue.lease (channel='C')
self.assertTrue(queue.delete (lease, 'DONE'))
self.assertFalse(queue.delete (lease, 'DONE'))
lease = queue.lease (channel='C')
self.assertTrue(queue.delete (lease, 'DONE'))
lease = queue.lease (channel='C')
self.assertIsNone(lease)
