Adds methods to store the locator map in the db

In order to use the db for system monitoring we need a way of mapping the locator the the uuids the appointments/jobs are identified in the db. First steps towards #15

pisa/db_manager.py

@@ -2,7 +2,8 @@ import json
 import plyvel
 
 from pisa.logger import Logger
-from pisa.conf import WATCHER_PREFIX, RESPONDER_PREFIX, WATCHER_LAST_BLOCK_KEY, RESPONDER_LAST_BLOCK_KEY
+from pisa.conf import WATCHER_PREFIX, RESPONDER_PREFIX, WATCHER_LAST_BLOCK_KEY, RESPONDER_LAST_BLOCK_KEY, \
+    LOCATOR_MAP_PREFIX
 
 logger = Logger("DBManager")
 
@@ -73,6 +74,40 @@ class DBManager:
         self.create_entry(uuid, job, prefix=RESPONDER_PREFIX)
         logger.info("Adding appointment to Responder's db", uuid=uuid)
 
+    def load_locator_map(self, locator):
+        key = (LOCATOR_MAP_PREFIX+locator).encode('utf-8')
+        locator_map = self.db.get(key)
+
+        if locator_map is not None:
+            locator_map = json.loads(locator_map.decode('utf-8'))
+
+        else:
+            logger.info("Locator not found in the db", locator=locator)
+
+        return locator_map
+
+    def store_update_locator_map(self, locator, uuid):
+        locator_map = self.load_locator_map(locator)
+
+        if locator_map is not None:
+            if uuid not in locator_map:
+                locator_map.append(uuid)
+                logger.info("Updating locator map", locator=locator, uuid=uuid)
+
+            else:
+                logger.info("UUID already in the map", locator=locator, uuid=uuid)
+
+        else:
+            locator_map = [uuid]
+            logger.info("Creating new locator map", locator=locator, uuid=uuid)
+
+        key = (LOCATOR_MAP_PREFIX + locator).encode('utf-8')
+        self.db.put(key, json.dumps(locator_map).encode('utf-8'))
+
+    def delete_locator_map(self, locator):
+        self.delete_entry(locator, prefix=LOCATOR_MAP_PREFIX)
+        logger.info("Deleting locator map from db", uuid=locator)
+
     def delete_watcher_appointment(self, uuid):
         self.delete_entry(uuid, prefix=WATCHER_PREFIX)
         logger.info("Deleting appointment from Watcher's db", uuid=uuid)

test/unit/test_db_manager.py

@@ -5,8 +5,8 @@ import shutil
 from uuid import uuid4
 
 from pisa.db_manager import DBManager
-from test.unit.conftest import get_random_value_hex, generate_dummy_appointment, generate_dummy_job
-from pisa.conf import WATCHER_LAST_BLOCK_KEY, RESPONDER_LAST_BLOCK_KEY
+from test.unit.conftest import get_random_value_hex, generate_dummy_appointment
+from pisa.conf import WATCHER_LAST_BLOCK_KEY, RESPONDER_LAST_BLOCK_KEY, LOCATOR_MAP_PREFIX
 
 
 @pytest.fixture(scope='module')
@@ -147,6 +147,41 @@ def test_load_responder_jobs_empty(db_manager):
     assert len(db_manager.load_responder_jobs()) == 0
 
 
+def test_load_locator_map_empty(db_manager):
+    assert db_manager.load_locator_map(get_random_value_hex(32)) is None
+
+
+def test_store_update_locator_map_empty(db_manager):
+    uuid = uuid4().hex
+    locator = get_random_value_hex(32)
+    db_manager.store_update_locator_map(locator, uuid)
+
+    # Check that the locator map has been properly stored
+    assert db_manager.load_locator_map(locator) == [uuid]
+
+    # If we try to add the same uuid again the list shouldn't change
+    db_manager.store_update_locator_map(locator, uuid)
+    assert db_manager.load_locator_map(locator) == [uuid]
+
+    # Add another uuid to the same locator and check that it also works
+    uuid2 = uuid4().hex
+    db_manager.store_update_locator_map(locator, uuid2)
+
+    assert set(db_manager.load_locator_map(locator)) == set([uuid, uuid2])
+
+
+def test_delete_locator_map(db_manager):
+    locator_maps = db_manager.load_appointments_db(prefix=LOCATOR_MAP_PREFIX)
+    assert(len(locator_maps) != 0)
+
+    for locator, uuids in locator_maps.items():
+        print(locator)
+        db_manager.delete_locator_map(locator)
+
+    locator_maps = db_manager.load_appointments_db(prefix=LOCATOR_MAP_PREFIX)
+    assert (len(locator_maps) == 0)
+
+
 def test_store_load_watcher_appointment(db_manager, watcher_appointments):
     for uuid, appointment in watcher_appointments.items():
         db_manager.store_watcher_appointment(uuid, appointment.to_json())