Adds end time to the user (subscription) information. It will replace the appointment one.

teos/gatekeeper.py

@@ -20,19 +20,48 @@ class IdentificationFailure(Exception):
     pass
 
 
+class UserInfo:
+    def __init__(self, available_slots, subscription_end_time, appointments=None):
+        self.available_slots = available_slots
+        self.subscription_end_time = subscription_end_time
+
+        if not appointments:
+            self.appointments = {}
+        else:
+            self.appointments = appointments
+
+    @classmethod
+    def from_dict(cls, user_data):
+        available_slots = user_data.get("available_slots")
+        appointments = user_data.get("appointments")
+        subscription_end_time = user_data.get("subscription_end_time")
+
+        if any(v is None for v in [available_slots, appointments, subscription_end_time]):
+            raise ValueError("Wrong appointment data, some fields are missing")
+
+        return cls(available_slots, subscription_expiry, appointments)
+
+    def to_dict(self):
+        return self.__dict__
+
+
 class Gatekeeper:
     """
     The :class:`Gatekeeper` is in charge of managing the access to the tower. Only registered users are allowed to
     perform actions.
 
     Attributes:
-        registered_users (:obj:`dict`): a map of user_pk:appointment_slots.
+        registered_users (:obj:`dict`): a map of user_pk:UserInfo.
     """
 
-    def __init__(self, user_db, default_slots):
+    def __init__(self, user_db, block_processor, default_slots, default_subscription_duration):
         self.default_slots = default_slots
+        self.block_processor = block_processor
+        self.default_subscription_duration = default_subscription_duration
         self.user_db = user_db
-        self.registered_users = user_db.load_all_users()
+        self.registered_users = {
+            user_id: UserInfo.from_dict(user_data) for user_id, user_data in user_db.load_all_users().items()
+        }
 
     def add_update_user(self, user_pk):
         """
@@ -42,20 +71,27 @@ class Gatekeeper:
             user_pk(:obj:`str`): the public key that identifies the user (33-bytes hex str).
 
         Returns:
-            :obj:`int`: the number of available slots in the user subscription.
+            :obj:`tuple`: a tuple with the number of available slots in the user subscription and the subscription end
+            time (in absolute block height).
         """
 
         if not is_compressed_pk(user_pk):
             raise ValueError("Provided public key does not match expected format (33-byte hex string)")
 
         if user_pk not in self.registered_users:
-            self.registered_users[user_pk] = {"available_slots": self.default_slots}
+            self.registered_users[user_pk] = UserInfo(
+                self.default_slots, self.block_processor.get_block_count() + self.default_subscription_duration
+            )
         else:
-            self.registered_users[user_pk]["available_slots"] += self.default_slots
+            # FIXME: For now new calls to register add default_slots to the current count and reset the expiry time
+            self.registered_users[user_pk].available_slots += self.default_slots
+            self.registered_users[user_pk].subscription_expiry = (
+                self.block_processor.get_block_count() + self.default_subscription_duration
+            )
 
-        self.user_db.store_user(user_pk, self.registered_users[user_pk])
+        self.user_db.store_user(user_pk, self.registered_users[user_pk].to_dict())
 
-        return self.registered_users[user_pk]["available_slots"]
+        return self.registered_users[user_pk].available_slots, self.registered_users[user_pk].subscription_end_time
 
     def identify_user(self, message, signature):
         """
@@ -97,9 +133,9 @@ class Gatekeeper:
         """
 
         # DISCUSS: we may want to return a different exception if the user does not exist
-        if user_pk in self.registered_users and n <= self.registered_users.get(user_pk).get("available_slots"):
-            self.registered_users[user_pk]["available_slots"] -= n
-            self.user_db.store_user(user_pk, self.registered_users[user_pk])
+        if user_pk in self.registered_users and n <= self.registered_users.get(user_pk).available_slots:
+            self.registered_users[user_pk].available_slots -= n
+            self.user_db.store_user(user_pk, self.registered_users[user_pk].to_dict())
         else:
             raise NotEnoughSlots(user_pk, n)
 
@@ -114,5 +150,5 @@ class Gatekeeper:
 
         # DISCUSS: if the user does not exist we may want to log or return an exception.
         if user_pk in self.registered_users:
-            self.registered_users[user_pk]["available_slots"] += n
-            self.user_db.store_user(user_pk, self.registered_users[user_pk])
+            self.registered_users[user_pk].available_slots += n
+            self.user_db.store_user(user_pk, self.registered_users[user_pk].to_dict())