diff --git a/teos/cleaner.py b/teos/cleaner.py index 0e64d67..938fcfd 100644 --- a/teos/cleaner.py +++ b/teos/cleaner.py @@ -238,3 +238,32 @@ class Cleaner: db_manager.batch_delete_responder_trackers(completed_trackers) db_manager.batch_delete_watcher_appointments(completed_trackers) db_manager.batch_delete_triggered_appointment_flag(completed_trackers) + + @staticmethod + def delete_gatekeeper_appointments(gatekeeper, appointment_to_delete): + """ + Deletes a list of expired / completed appointments of a given user both from memory and the UserDB. + + Args: + gatekeeper (:obj:`Gatekeeper `): a `Gatekeeper` instance in charge to control + the user access and subscription expiry. + appointment_to_delete (:obj:`dict`): uuid:user_id dict containing the appointments to delete + (expired + completed) + """ + + user_ids = [] + # Remove appointments from memory + for uuid, user_id in appointment_to_delete.items(): + if user_id in gatekeeper.registered_users and uuid in gatekeeper.registered_users[user_id].appointments: + # Remove the appointment from the appointment list and update the available slots + gatekeeper.lock.acquire() + freed_slots = gatekeeper.registered_users[user_id].appointments.pop(uuid) + gatekeeper.registered_users[user_id].available_slots += freed_slots + gatekeeper.lock.release() + + if user_id not in user_ids: + user_ids.append(user_id) + + # Store the updated users in the DB + for user_id in user_ids: + gatekeeper.user_db.store_user(user_id, gatekeeper.registered_users[user_id].to_dict()) diff --git a/test/teos/unit/test_cleaner.py b/test/teos/unit/test_cleaner.py index 163dde0..5ea02a1 100644 --- a/test/teos/unit/test_cleaner.py +++ b/test/teos/unit/test_cleaner.py @@ -1,8 +1,9 @@ import random from uuid import uuid4 -from teos.responder import TransactionTracker from teos.cleaner import Cleaner +from teos.gatekeeper import UserInfo +from teos.responder import TransactionTracker from common.appointment import Appointment from test.teos.unit.conftest import get_random_value_hex @@ -206,3 +207,36 @@ def test_delete_trackers_no_db_match(db_manager): # We should be able to delete the correct ones and not fail in the others Cleaner.delete_trackers(completed_trackers, height, trackers, tx_tracker_map, db_manager) assert not set(completed_trackers).issubset(trackers.keys()) + + +def test_delete_gatekeeper_appointments(gatekeeper): + # delete_gatekeeper_appointments should delete the appointments from user as long as both exist + + appointments_not_to_delete = {} + appointments_to_delete = {} + # Let's add some users and appointments to the Gatekeeper + for _ in range(10): + user_id = get_random_value_hex(16) + # The UserInfo params do not matter much here + gatekeeper.registered_users[user_id] = UserInfo(available_slots=100, subscription_expiry=0) + for _ in range(random.randint(0, 10)): + # Add some appointments + uuid = get_random_value_hex(16) + gatekeeper.registered_users[user_id].appointments[uuid] = 1 + + if random.randint(0, 1) % 2: + appointments_to_delete[uuid] = user_id + else: + appointments_not_to_delete[uuid] = user_id + + # Now let's delete half of them + Cleaner.delete_gatekeeper_appointments(gatekeeper, appointments_to_delete) + + all_appointments_gatekeeper = [] + # Let's get all the appointments in the Gatekeeper + for user_id, user in gatekeeper.registered_users.items(): + all_appointments_gatekeeper.extend(user.appointments) + + # Check that the first half of the appointments are not in the Gatekeeper, but the second half is + assert not set(appointments_to_delete).issubset(all_appointments_gatekeeper) + assert set(appointments_not_to_delete).issubset(all_appointments_gatekeeper)