SIENTIAPDE-988 Implement core functionality for OPC ingestor, including main execution flow, IngestorManager methods, and OPC server management. Add unit tests for IngestorManager and OPCManager.

ingestor/ingestor.py

@@ -0,0 +1,76 @@
+from logging import Formatter, StreamHandler, getLogger
+from ingestor.managers.ingestor_manager import IngestorManager
+from os import getenv
+from time import sleep
+
+
+def main():
+    # Get os parameters
+    kafka_servers = getenv("KAFKA_SERVERS")
+    redis_host = getenv("REDIS_HOST")
+    redis_port = int(getenv("REDIS_PORT"))
+    lease_ttl = int(getenv("LEASE_TTL"))
+    heartbeat_ttl = int(getenv("HEARTBEAT_TTL"))
+    pod_id = getenv("HOSTNAME")
+    number_of_ingestors = int(getenv("REPLICA_COUNT"))
+    poll_interval = int(getenv("POLL_INTERVAL"))
+
+    logger = getLogger(__name__)
+    logger.setLevel(getenv("LOG_LEVEL", "INFO"))
+    handler = StreamHandler()
+    formatter = Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    handler.setFormatter(formatter)
+
+    logger.addHandler(handler)
+
+    ingestor_manager = IngestorManager(
+        kafka_servers, redis_host, redis_port,
+        lease_ttl, heartbeat_ttl, pod_id,
+        number_of_ingestors, poll_interval, logger
+    )
+    # Declare ingestor ative
+    ingestor_manager.declare_active()
+    # Get slot lease
+    acquired = ingestor_manager.get_slot_leases()
+
+    # Subscribe to acquired slots
+    ingestor_manager.update_opc_servers()
+    ingestor_manager.subscribe_to_tags(acquired)
+
+    while True:
+        # Declare ingestor as active
+        ingestor_manager.declare_active()
+
+        # Update opc servers
+        ingestor_manager.update_slot_config()
+
+        # Get active ingestors
+        ingestors = ingestor_manager.get_active_ingestors()
+
+        ingestor_diff = number_of_ingestors - len(ingestors)
+        slot_diff = len(ingestor_manager.managed_tags) - 1
+
+        if ingestor_diff > 0:
+            # Some ingestors are innactive, so theres "ingestor_diff" slots available
+
+            # Get slot lease
+            acquired = ingestor_manager.get_slot_leases(ingestor_diff)
+
+            # Subscribe to acquired slots
+            ingestor_manager.update_opc_servers()
+            ingestor_manager.subscribe_to_tags(acquired)
+
+        elif slot_diff > 0:
+            # Some ingestors are active and without slots, so we need to drop
+
+            overleases = list(ingestor_manager.managed_tags.keys())[1:]
+
+            ingestor_manager.drop_slot_leases(overleases)
+
+        # Sleep for poll interval
+        sleep(poll_interval)
+
+
+if __name__ == "__main__":
+    main()

ingestor/managers/ingestor_manager.py

@@ -8,29 +8,96 @@ from ingestor.managers.resource_manager import ResourceManager
 class IngestorManager():
     def __init__(self,
                  kafka_servers: str, redis_host: str, redis_port: int,
-                 lease_ttl: int, heartbeat_ttl: int, poll_interval: int,
+                 lease_ttl: int, heartbeat_ttl: int,  pod_id: str,
-                 logger: Logger):
+                 number_of_ingestors: int,
+                 poll_interval: int, logger: Logger):
 
         self.data_manager = DataManager(kafka_servers, logger)
-        self.opc_managers = []
+        self.opc_managers = {}
         self.resource_manager = ResourceManager(
-            redis_host, redis_port, lease_ttl, heartbeat_ttl, logger
+            redis_host, redis_port, lease_ttl, heartbeat_ttl, pod_id
         )
+        self.number_of_ingestors = number_of_ingestors
+        self.poll_interval = poll_interval
+        self.logger = logger
+        self.managed_tags = {}
+        self.opc_servers = {}
 
-    def init_ingestor(self):
+    def update_opc_servers(self):
-        pass
+        for slot, slot_config in self.managed_tags.items():
+            for server, server_config in slot_config.items():
+                server_config = server_config.copy()
+                server_config.pop('tags', None)
+                if server not in self.opc_servers:
+                    self.opc_managers[server] = OpcManager.initialize_from_config(
+                        server_config, self.data_manager, self.logger
+                    )
+                    self.opc_managers[server].connect()
+                elif self.opc_servers[server] != server_config:
+                    del self.opc_managers[server]
+                    self.opc_managers[server] = OpcManager.initialize_from_config(
+                        server_config, self.data_manager, self.logger
+                    )
+                    self.opc_managers[server].connect()
+
+                self.opc_servers[server] = server_config
 
     def declare_active(self):
-        pass
+        self.resource_manager.ingestor_heartbeat()
 
-    def get_active_ingestors(self):
+    def get_active_ingestors(self) -> List[str]:
-        pass
+        self.resource_manager.get_all_ingestors()
 
-    def get_slot_leases(self, max_slots: int = 1) -> List[Dict]:
+    def get_slot_leases(self, max_slots: int = 1) -> Dict:
-        pass
+        acquired = {}
+        for i in range(1, self.number_of_ingestors + 1):
+            if self.resource_manager.lease_tag(str(i)):
+                self.logger.info(f"Leased slot {i}")
+                acquired[str(i)] = self.resource_manager.get_tag_slot(str(i))
+
+                if len(acquired) >= max_slots:
+                    self.managed_tags.update(acquired)
+                    return acquired
+
+        self.logger.warning(
+            f"Unable to acquire {max_slots} slots. "
+            f"Only {acquired} slots were leased."
+        )
+        self.managed_tags.update(acquired)
+        return acquired
+
+    def update_slot_config(self) -> Dict:
+        for slot, slot_config in self.managed_tags.items():
+            update = self.resource_manager.get_tag_slot(slot)
+
+            if update != slot_config:
+                self.logger.info(
+                    f"Slot {slot} configuration updated. "
+                    f"Old: {slot_config}, New: {update}"
+                )
+                self.managed_tags[slot] = update
+
+                self.update_opc_servers()
+                self.subscribe_to_tags(update)
+
+        return update
 
     def drop_slot_leases(self, ids: List[str]) -> None:
-        pass
+        for id in ids:
+            self.resource_manager.drop_tag_lease(id)
 
-    def subscribe_to_tags(self, tags: List[Dict]) -> None:
+    def subscribe_to_tags(self, tags: Dict) -> None:
-        pass
+        for slot, slot_config in tags.items():
+            for server, server_config in slot_config.items():
+                if server not in self.opc_managers:
+                    self.logger.error(
+                        f"Server {server} not found in opc_managers."
+                    )
+                    continue
+                if slot not in self.opc_managers[server].subscriptions:
+                    self.opc_managers[server].create_subscription(
+                        slot
+                    )
+                self.opc_managers[server].subscribe(
+                    slot, server_config['tags'], self.poll_interval
+                )

ingestor/managers/opc_manager.py

@@ -24,9 +24,23 @@ class OpcManager():
         self.private_key_path = private_key_path
         self.server_cert_path = server_cert_path
         self.nodes = {}
-        self.subscription = None
+        self.subscriptions = {}
         self.data_manager = data_manager
 
+    def initialize_from_config(self, server_config: dict, data_manager: DataManager, logger: Logger):
+        """
+        Initializes the OpcManager instance using a server configuration dictionary.
+        Args:
+            server_config (dict): A dictionary containing server configuration details.
+            data_manager (DataManager): An instance of DataManager for data handling.
+            logger (Logger): Logger instance for logging information.
+        """
+        self.__init__(
+            server_config['name'], server_config['url'], data_manager, logger, server_config['server_uri'],
+            server_config.get('cert_path'), server_config.get(
+                'private_key_path'), server_config.get('server_cert_path')
+        )
+
     def set_security(self):
         """
         Configures the security settings for the OPC UA client.
@@ -82,7 +96,7 @@ class OpcManager():
         self.logger.info('Starting connection...')
         self.client.connect()
 
-    def create_subscription(self, period: int = 500):
+    def create_subscription(self, name: str, period: int = 500):
         """
         Creates a subscription with the specified monitoring period.
         This method establishes a subscription to monitor data changes or events
@@ -101,12 +115,11 @@ class OpcManager():
             raise ValueError("Client not connected. Call connect first.")
 
         p = period if period != None else 500
-        self.period = p
+        self.subscriptions[name] = self.client.create_subscription(
-        self.subscription = self.client.create_subscription(
             p, self)
         self.logger.info('Subscription created.')
 
-    def subscribe(self, nodes: dict, collect_period: int):
+    def subscribe(self, subscription: str, nodes: dict, collect_period: int):
         """
         Subscribes to a set of OPC UA nodes for data change notifications.
         This method adds the specified nodes to the subscription and configures
@@ -123,7 +136,7 @@ class OpcManager():
                 `create_subscription` prior to this method.
         """
 
-        if not self.subscription:
+        if not self.subscriptions.get(subscription):
             raise ValueError(
                 "Subscription not created. Call create_subscription first.")
 
@@ -138,7 +151,16 @@ class OpcManager():
                 'cycle_count': 0
             }
 
-        self.subscription.subscribe_data_change(self.addr_nodes)
+        self.subscriptions[subscription].subscribe_data_change(self.addr_nodes)
+
+    def unsubscribe(self, subscription: str):
+
+        if not self.subscriptions.get(subscription):
+            raise ValueError(
+                "Subscription not created. Call create_subscription first.")
+        self.subscriptions[subscription].delete()
+        del self.subscriptions[subscription]
+        self.logger.info(f"Unsubscribed from {subscription}.")
 
     def __del__(self):
         self.disconnect()
@@ -156,33 +178,13 @@ class OpcManager():
 
         self.logger.warning('Disconnecting from OPC server')
         try:
-            self.subscription.delete()
+            [self.subscriptions[sub].delete() for sub in self.subscriptions]
+            self.logger.warning("Deleted all subscriptions.")
             self.client.disconnect()
             self.logger.warning("Disconnected from OPC UA server.")
         except Exception as sub_error:
             self.logger.error(f"Failed to clean up subscription: {sub_error}")
 
-    def init_collector(self, nodes: dict, collect_period: int, period: int):
-        """
-        Initializes the OPC data collector by connecting to the server, creating a subscription,
-        and subscribing to the specified nodes.
-        Args:
-            nodes (dict): A dictionary where keys are node identifiers and values are the node details
-                          to be subscribed to.
-            collect_period (int): The interval (in milliseconds) at which data should be collected
-                                  from the subscribed nodes.
-            period (int): The publishing interval (in milliseconds) for the subscription.
-        Returns:
-            None
-        """
-
-        self.connect()
-        self.create_subscription(period)
-        self.subscribe(nodes, collect_period)
-
-        self.logger.info(
-            f"Connected to OPC server {self.name} at {self.url}.")
-
     def datachange_notification(self, node, _val, data):
         """
         Handles data change notifications for monitored OPC UA nodes.
@@ -263,7 +265,7 @@ class OpcManager():
         Monitors the OPC connection and triggers events based on the number of cycles
         without receiving data from the OPC server.
         Args:
-            handle_listen_events (Callable[[str, str, NotificationLevel], None]): 
+            handle_listen_events (Callable[[str, str, NotificationLevel], None]):
                 A callback function to handle notification events. It takes three arguments:
                 - A message string describing the event.
                 - An event code string.

ingestor/managers/resource_manager.py

@@ -1,4 +1,5 @@
 import json
+from typing import List
 from redis import Redis
 
 
@@ -48,19 +49,19 @@ class ResourceManager:
         self.redis.set(
             f"heartbeat:ingestor:{self.pod_id}", 1, ex=self.heartbeat_ttl)
 
-    def lease_tag(self, tag_id: str) -> None:
+    def lease_tag(self, tag_id: str) -> bool:
         """
-        Acquires a lease for a specific OPC tag by setting a key in Redis with a 
+        Attempts to lease a tag by setting a key in Redis with a specified TTL (time-to-live).
-        time-to-live (TTL). This ensures that the tag is associated with the 
+        This method uses the Redis `SET` command with the `NX` option to ensure that the key 
-        current pod for a limited duration.
+        is only set if it does not already exist. The key is set with an expiration time 
+        defined by `lease_ttl`.
         Args:
-            tag_id (str): The unique identifier of the OPC tag to lease.
+            tag_id (str): The unique identifier of the tag to be leased.
-        Raises:
+        Returns:
-            redis.exceptions.RedisError: If there is an issue communicating with 
+            bool: True if the lease was successfully acquired, False otherwise.
-                the Redis server.
         """
 
-        self.redis.set(
+        return self.redis.set(
             f"lease:opc_tags:{tag_id}", self.pod_id, nx=True, ex=self.lease_ttl)
 
     def renew_tag_lease(self, tag_id: str) -> bool:
@@ -82,3 +83,26 @@ class ResourceManager:
             return True
 
         return False
+
+    def drop_tag_lease(self, tag_id: str) -> None:
+        """
+        Drops the lease for a specific OPC tag.
+        This method removes the lease for the given OPC tag by deleting the corresponding key in Redis.
+        Args:
+            tag_id (str): The identifier of the OPC tag whose lease is to be dropped.
+        Returns:
+            None
+        """
+
+        self.redis.delete(f"lease:opc_tags:{tag_id}")
+
+    def get_all_ingestors(self) -> List[str]:
+        """
+        Retrieves all active ingestors from Redis.
+        This method fetches all keys in Redis that match the pattern for ingestor leases
+        and returns a list of active ingestors.
+        Returns:
+            list: A list of active ingestors.
+        """
+
+        return self.redis.keys("heartbeat:ingestor:*")

tests/managers/test_ingestor_manager.py

@@ -0,0 +1,227 @@
+from unittest.mock import MagicMock, patch
+from pytest import fixture
+
+from ingestor.managers.ingestor_manager import IngestorManager
+
+
+@fixture
+@patch('ingestor.managers.ingestor_manager.DataManager')
+@patch('ingestor.managers.ingestor_manager.ResourceManager')
+def ingestor_manager(data_manager_mock, resource_manager_mock):
+    return IngestorManager(
+        kafka_servers="localhost:9092",
+        redis_host="localhost",
+        redis_port=6379,
+        lease_ttl=60,
+        heartbeat_ttl=60,
+        pod_id="test_pod",
+        number_of_ingestors=3,
+        poll_interval=5,
+        logger=MagicMock()
+    )
+
+
+@patch('ingestor.managers.ingestor_manager.OpcManager')
+@patch('ingestor.managers.ingestor_manager.DataManager')
+@patch('ingestor.managers.ingestor_manager.ResourceManager')
+def test___init__(resource_manager_mock, data_manager_mock, opc_manager_mock):
+
+    ingestor = IngestorManager(
+        kafka_servers="localhost:9092",
+        redis_host="localhost",
+        redis_port=6379,
+        lease_ttl=60,
+        heartbeat_ttl=60,
+        pod_id="test_pod",
+        number_of_ingestors=3,
+        poll_interval=5,
+        logger=MagicMock()
+    )
+
+    opc_manager_mock.assert_not_called()
+    data_manager_mock.assert_called_once_with(
+        "localhost:9092", ingestor.logger)
+    resource_manager_mock.assert_called_once_with(
+        "localhost", 6379, 60, 60, "test_pod")
+
+    assert ingestor.number_of_ingestors == 3
+    assert ingestor.poll_interval == 5
+    assert ingestor.managed_tags == {}
+    assert ingestor.opc_servers == {}
+    assert ingestor.opc_managers == {}
+    assert ingestor.data_manager == data_manager_mock.return_value
+    assert ingestor.resource_manager == resource_manager_mock.return_value
+
+
+@patch('ingestor.managers.ingestor_manager.OpcManager')
+def test_update_opc_servers(opc_manager, ingestor_manager):
+    manager1 = MagicMock()
+    manager2 = MagicMock()
+    manager3 = MagicMock()
+
+    def mock_initialize_from_config(config, data_manager, logger):
+        if config == {"config": "config1"}:
+            return manager1
+        elif config == {"config": "config2"}:
+            return manager2
+        elif config == {"config": "config3"}:
+            return manager3
+
+    opc_manager.initialize_from_config = MagicMock(
+        side_effect=mock_initialize_from_config
+    )
+
+    ingestor_manager.managed_tags = {
+        "slot1": {
+            "server1": {"config": "config1"},
+            "server2": {"config": "config2"}
+        },
+        "slot2": {
+            "server3": {"config": "config3"},
+            "server1": {"config": "config1"}
+        }
+    }
+
+    ingestor_manager.opc_servers = {
+        "server3": {"config": "config3"},
+        "server2": {"old_config": "old_config2"}
+    }
+
+    mock = MagicMock()
+    ingestor_manager.opc_managers['server3'] = MagicMock()
+    ingestor_manager.opc_managers['server2'] = mock
+    ingestor_manager.update_opc_servers()
+
+    assert len(ingestor_manager.opc_managers) == 3
+    assert len(ingestor_manager.opc_servers) == 3
+
+    opc_manager.initialize_from_config.assert_any_call(
+        {"config": "config1"}, ingestor_manager.data_manager, ingestor_manager.logger)
+    opc_manager.initialize_from_config.assert_any_call(
+        {"config": "config2"}, ingestor_manager.data_manager, ingestor_manager.logger)
+
+    ingestor_manager.opc_managers['server1'].connect.assert_called_once()
+    ingestor_manager.opc_managers['server2'].connect.assert_called_once()
+    ingestor_manager.opc_managers['server3'].connect.assert_not_called()
+
+    assert ingestor_manager.opc_servers['server1'] == {"config": "config1"}
+    assert ingestor_manager.opc_servers['server2'] == {"config": "config2"}
+    assert ingestor_manager.opc_servers['server3'] == {"config": "config3"}
+
+    assert ingestor_manager.opc_managers['server2'] != mock
+
+
+def test_declare_active(ingestor_manager):
+    ingestor_manager.resource_manager.ingestor_heartbeat = MagicMock()
+    ingestor_manager.declare_active()
+    ingestor_manager.resource_manager.ingestor_heartbeat.assert_called_once()
+
+
+def test_get_active_ingestors(ingestor_manager):
+    ingestor_manager.resource_manager.get_all_ingestors = MagicMock()
+    ingestor_manager.get_active_ingestors()
+    ingestor_manager.resource_manager.get_all_ingestors.assert_called_once()
+
+
+def test_get_slot_leases_1_success(ingestor_manager):
+    ingestor_manager.resource_manager.lease_tag = MagicMock(
+        return_value=True)
+    ingestor_manager.resource_manager.get_tag_slot = MagicMock(
+        return_value={"tags": ["tag1"]})
+
+    result = ingestor_manager.get_slot_leases()
+
+    assert result == {
+        "1": {"tags": ["tag1"]}
+    }
+
+
+def test_get_slot_leases_2_success(ingestor_manager):
+    ingestor_manager.resource_manager.lease_tag = MagicMock(
+        side_effect=[True, True])
+    ingestor_manager.resource_manager.get_tag_slot = MagicMock(
+        side_effect=[{"tags": ["tag1"]}, {"tags": ["tag2"]}])
+
+    result = ingestor_manager.get_slot_leases(max_slots=2)
+
+    assert result == {
+        "1": {"tags": ["tag1"]},
+        "2": {"tags": ["tag2"]}
+    }
+
+
+def test_get_slot_leases_1_failure(ingestor_manager):
+    ingestor_manager.resource_manager.lease_tag = MagicMock(
+        return_value=False)
+    ingestor_manager.resource_manager.get_tag_slot = MagicMock(
+        return_value={"tags": ["tag1"]})
+
+    result = ingestor_manager.get_slot_leases()
+    ingestor_manager.resource_manager.get_tag_slot.assert_not_called()
+
+    assert result == {}
+
+
+def test_update_slot_config(ingestor_manager):
+    ingestor_manager.managed_tags = {
+        "slot1": {"config": "old_config"},
+        "slot2": {"config": "new_config"}
+    }
+
+    ingestor_manager.resource_manager.get_tag_slot = MagicMock(
+        side_effect=[
+            {"config": "updated_config"},
+            {"config": "new_config"}
+        ]
+    )
+
+    ingestor_manager.update_opc_servers = MagicMock()
+    ingestor_manager.subscribe_to_tags = MagicMock()
+
+    ingestor_manager.update_slot_config()
+
+    assert ingestor_manager.managed_tags["slot1"] == {
+        "config": "updated_config"}
+    assert ingestor_manager.managed_tags["slot2"] == {
+        "config": "new_config"}
+
+    ingestor_manager.update_opc_servers.assert_called_once()
+    ingestor_manager.subscribe_to_tags.assert_called_once_with(
+        {"config": "updated_config"}
+    )
+
+
+def test_drop_slot_leases(ingestor_manager):
+    ingestor_manager.resource_manager.drop_tag_lease = MagicMock()
+    ingestor_manager.drop_slot_leases(["1", "2"])
+
+    ingestor_manager.resource_manager.drop_tag_lease.assert_any_call("1")
+    ingestor_manager.resource_manager.drop_tag_lease.assert_any_call("2")
+
+
+def test_subscribe_to_tags(ingestor_manager):
+    ingestor_manager.opc_managers = {
+        "server1": MagicMock(),
+        "server2": MagicMock()
+    }
+    ingestor_manager.subscriptions = {
+        "server1": MagicMock()
+    }
+    tags = {
+        'slot1': {
+            "server1": {"tags": "config1"},
+            "server2": {"tags": "config2"},
+            'server3': {"tags": "config3"}
+        }
+    }
+
+    ingestor_manager.subscribe_to_tags(tags)
+
+    ingestor_manager.opc_managers["server1"].create_subscription.assert_called_once_with(
+        'slot1')
+
+    ingestor_manager.opc_managers["server1"].subscribe.assert_called_once_with(
+        'slot1', 'config1', ingestor_manager.poll_interval)
+    ingestor_manager.opc_managers["server2"].subscribe.assert_called_once_with(
+        'slot1', 'config2', ingestor_manager.poll_interval)
+    ingestor_manager.opc_managers.get("server3") is None

tests/managers/test_opc_manager.py

@@ -51,7 +51,7 @@ def opc_manager(raw_opc_manager):
 
 @fixture
 def opc_manager_subscribed(opc_manager):
-    opc_manager.subscription = MagicMock()
+    opc_manager.subscriptions['sub1'] = MagicMock()
 
     return opc_manager
 
@@ -114,7 +114,7 @@ def test_connect_with_security(client, raw_opc_manager):
 
 def test_create_subscription_no_client(raw_opc_manager):
     try:
-        raw_opc_manager.create_subscription()
+        raw_opc_manager.create_subscription('sub1')
     except ValueError as e:
         assert str(e) == "Client not connected. Call connect first."
     else:
@@ -122,26 +122,24 @@ def test_create_subscription_no_client(raw_opc_manager):
 
 
 def test_create_subscription_success_has_period(opc_manager):
-    opc_manager.create_subscription(1000)
+    opc_manager.create_subscription('sub1', 1000)
 
-    assert opc_manager.period == 1000
     opc_manager.client.create_subscription.assert_called_once_with(
         1000, opc_manager)
-    assert opc_manager.subscription is not None
+    assert opc_manager.subscriptions['sub1'] is not None
 
 
 def test_create_subscription_success_no_period(opc_manager):
-    opc_manager.create_subscription(None)
+    opc_manager.create_subscription('sub1', None)
 
-    assert opc_manager.period == 500
     opc_manager.client.create_subscription.assert_called_once_with(
         500, opc_manager)
-    assert opc_manager.subscription is not None
+    assert opc_manager.subscriptions['sub1'] is not None
 
 
 def test_subscribe_no_subscription(opc_manager):
     try:
-        opc_manager.subscribe(tags, 1000)
+        opc_manager.subscribe('sub1', tags, 1000)
     except ValueError as e:
         assert str(
             e) == "Subscription not created. Call create_subscription first."
@@ -160,42 +158,45 @@ def test_subscribe_success(opc_manager_subscribed):
         opc_manager_subscribed.client.get_node(n) for n in tags if n != 'ns=3;i=1001']
 
 
+def test_unsubscribe_no_subscription(opc_manager):
+    try:
+        opc_manager.unsubscribe('sub1')
+    except ValueError as e:
+        assert str(
+            e) == "Subscription not created. Call create_subscription first."
+    else:
+        assert False, "ValueError not raised"
+
+
+def test_unsubscribe_success(opc_manager_subscribed):
+    opc_manager_subscribed.unsubscribe('sub1')
+
+    opc_manager_subscribed.subscriptions.get('sub1') is None
+
+
 def test_disconnect_success(opc_manager_subscribed):
     opc_manager_subscribed.client = MagicMock()
-    opc_manager_subscribed.subscription = MagicMock()
 
     opc_manager_subscribed.disconnect()
 
-    opc_manager_subscribed.subscription.delete.assert_called_once()
+    opc_manager_subscribed.subscriptions['sub1'].delete.assert_called_once()
     opc_manager_subscribed.client.disconnect.assert_called_once()
 
 
 def test_disconnect_error(opc_manager_subscribed):
     opc_manager_subscribed.client = MagicMock()
-    opc_manager_subscribed.subscription = MagicMock(
+    opc_manager_subscribed.subscriptions['sub1'] = MagicMock(
         delete=MagicMock(side_effect=Exception("Test error"))
     )
 
     opc_manager_subscribed.disconnect()
 
-    opc_manager_subscribed.subscription.delete.assert_called_once()
+    opc_manager_subscribed.subscriptions['sub1'].delete.assert_called_once()
     opc_manager_subscribed.client.disconnect.assert_not_called()
     opc_manager_subscribed.logger.error.assert_called_once_with(
         "Failed to clean up subscription: Test error")
 
 
-def test_init_collector(opc_manager_subscribed):
-    opc_manager_subscribed.connect = MagicMock()
-    opc_manager_subscribed.create_subscription = MagicMock()
-    opc_manager_subscribed.subscribe = MagicMock()
-
-    opc_manager_subscribed.init_collector(tags, 1000, 500)
-
-    opc_manager_subscribed.connect.assert_called_once()
-    opc_manager_subscribed.create_subscription.assert_called_once_with(500)
-    opc_manager_subscribed.subscribe.assert_called_once_with(tags, 1000)
-
-
 def test_datachange_notification(opc_manager_subscribed):
     data = MagicMock(
         monitored_item=MagicMock(

tests/managers/test_resource_manager.py

@@ -43,7 +43,8 @@ def test_ingestor_heartbeat(resource_manager):
 
 def test_lease_tag(resource_manager):
     resource_manager.redis.set.return_value = True
-    resource_manager.lease_tag('tag_id')
+    output = resource_manager.lease_tag('tag_id')
+    assert output is True
     resource_manager.redis.set.assert_called_once_with(
         'lease:opc_tags:tag_id', 'pod_id', nx=True, ex=10
     )
@@ -71,3 +72,11 @@ def test_renew_tag_lease_failure(resource_manager):
         'lease:opc_tags:tag_id'
     )
     resource_manager.redis.expire.assert_not_called()
+
+
+def test_drop_tag_lease(resource_manager):
+    resource_manager.redis.delete.return_value = True
+    resource_manager.drop_tag_lease('tag_id')
+    resource_manager.redis.delete.assert_called_once_with(
+        'lease:opc_tags:tag_id'
+    )