diff --git a/inventory_provider/tasks/common.py b/inventory_provider/tasks/common.py new file mode 100644 index 0000000000000000000000000000000000000000..6868ffa33e95752119070b0f00eea74be13e04a9 --- /dev/null +++ b/inventory_provider/tasks/common.py @@ -0,0 +1,8 @@ +import redis + + +def get_redis(config): + return redis.StrictRedis( + host=config['redis']['hostname'], + port=config['redis']['port']) + diff --git a/inventory_provider/tasks/worker.py b/inventory_provider/tasks/worker.py index 9923dc9430196fb8141e5085d38c5750fd5d833e..50a202a427b38967a266e4836e1320059d194028 100644 --- a/inventory_provider/tasks/worker.py +++ b/inventory_provider/tasks/worker.py @@ -3,10 +3,10 @@ import logging from celery import bootsteps, Task from collections import defaultdict -import redis from lxml import etree from inventory_provider.tasks.app import app +from inventory_provider.tasks.common import get_redis from inventory_provider import config from inventory_provider import constants from inventory_provider import environment @@ -29,9 +29,7 @@ class InventoryTask(Task): def save_key(hostname, key, value): assert isinstance(value, str), \ "sanity failure: expected string data as value" - r = redis.StrictRedis( - host=InventoryTask.config["redis"]["hostname"], - port=InventoryTask.config["redis"]["port"]) + r = get_redis(InventoryTask.config) r.hset( name=hostname, key=key, @@ -44,9 +42,7 @@ class InventoryTask(Task): def save_value(key, value): assert isinstance(value, str), \ "sanity failure: expected string data as value" - r = redis.StrictRedis( - host=InventoryTask.config["redis"]["hostname"], - port=InventoryTask.config["redis"]["port"]) + r = get_redis(InventoryTask.config) r.set( name=key, value=value) @@ -147,9 +143,7 @@ def update_alarmsdb_cache(self): @app.task() def update_interfaces_to_services(): # todo - factor this connection stuff out - r = redis.StrictRedis( - host=InventoryTask.config["redis"]["hostname"], - port=InventoryTask.config["redis"]["port"]) + r = get_redis(InventoryTask.config) with db.connection(InventoryTask.config["ops-db"]) as cx: services = opsdb.get_circuits(cx) @@ -177,9 +171,7 @@ def update_interfaces_to_services(): @app.task() def update_equipment_locations(): # todo - factor this connection stuff out - r = redis.StrictRedis( - host=InventoryTask.config["redis"]["hostname"], - port=InventoryTask.config["redis"]["port"]) + r = get_redis(InventoryTask.config) r.delete(equipment_locations_key) with db.connection(InventoryTask.config["ops-db"]) as cx: @@ -191,9 +183,7 @@ def update_equipment_locations(): @app.task() def update_circuit_hierarchy(): # todo - factor this connection stuff out - r = redis.StrictRedis( - host=InventoryTask.config["redis"]["hostname"], - port=InventoryTask.config["redis"]["port"]) + r = get_redis(InventoryTask.config) children_to_parents = defaultdict(list) parents_to_children = defaultdict(list) with db.connection(InventoryTask.config["ops-db"]) as cx: @@ -217,9 +207,7 @@ def update_circuit_hierarchy(): @app.task() def update_interface_statuses(): # todo - factor this connection stuff out - r = redis.StrictRedis( - host=InventoryTask.config["redis"]["hostname"], - port=InventoryTask.config["redis"]["port"]) + r = get_redis(InventoryTask.config) with db.connection(InventoryTask.config["ops-db"]) as cx: services = opsdb.get_circuits(cx) with db.connection(InventoryTask.config["alarms-db"]) as cx: diff --git a/test/test_celery_worker.py b/test/test_celery_worker.py index 64b57851949242dcc66f4e1c4c5a77ce4723c8b2..faa76da0b0fc4cf59778ad84af702e21d7f5c260 100644 --- a/test/test_celery_worker.py +++ b/test/test_celery_worker.py @@ -36,7 +36,7 @@ def mocked_worker_module(mocker, data_config): MockedRedis.db = {} mocker.patch( - 'inventory_provider.tasks.worker.redis.StrictRedis', + 'inventory_provider.tasks.common.redis.StrictRedis', MockedRedis)