From 6a9a4afd89f1d8243ea3fd81aa902b2f21ee7ec6 Mon Sep 17 00:00:00 2001
From: MrKevinWeiss <weiss.kevin604@gmail.com>
Date: Thu, 22 Feb 2024 11:48:46 +0100
Subject: [PATCH] feat(update_cache): Add power control to cache

---
 src/inet_nm/cli_update_cache.py | 15 ++++++++--
 src/inet_nm/location.py         | 33 ++++++++++++++++++++
 tests/test_location.py          | 53 +++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_location.py

diff --git a/src/inet_nm/cli_update_cache.py b/src/inet_nm/cli_update_cache.py
index 0114d2d..9f70621 100644
--- a/src/inet_nm/cli_update_cache.py
+++ b/src/inet_nm/cli_update_cache.py
@@ -3,6 +3,7 @@
 import inet_nm.config as cfg
 import inet_nm.location as loc
 from inet_nm._helpers import nm_print
+from inet_nm.power_control import DEFAULT_MAX_ALLOWED_NODES, PowerControl
 
 
 def _main():
@@ -14,9 +15,17 @@ def _main():
     nodes = cfg.NodesConfig(config_dir=args.config).load()
     loc_cache = cfg.LocationCache(config_dir=args.config)
     loc_cache.check_file(writable=True)
-
-    cache = loc.get_location_cache(nodes, loc_mapping)
-
+    caches = []
+    with PowerControl(
+        locations=loc_mapping,
+        nodes=nodes,
+        max_powered_devices=DEFAULT_MAX_ALLOWED_NODES,
+    ) as pc:
+        while not pc.power_on_complete:
+            pc.power_on_chunk()
+            caches.append(loc.get_location_cache(nodes, loc_mapping))
+            pc.power_off_unused()
+    cache = loc.merge_location_cache_chunks(caches)
     loc_cache.save(cache)
     nm_print(f"Updated {loc_cache.file_path}")
 
diff --git a/src/inet_nm/location.py b/src/inet_nm/location.py
index da61eb8..7bf0fc6 100644
--- a/src/inet_nm/location.py
+++ b/src/inet_nm/location.py
@@ -4,6 +4,39 @@
 from inet_nm.data_types import NmNode
 
 
+def merge_location_cache_chunks(caches: List[List[Dict]]):
+    """
+    Merge location cache chunks into a single cache.
+
+    Due to only being able to power on a chunk at a time we need to sort
+    through each of the location caches and look through all id_paths that
+    have a missing state and see if they are available in another chunk.
+    If they are then probably they were just powered off.
+
+    Args:
+        caches: List of location cache chunks.
+
+    Returns:
+        The merged location cache.
+    """
+    # TODO: Also check if all id_paths that are attached have the same node_uid
+    tmp_cache = {}
+    for chunk in caches:
+        for entry in chunk:
+            # If entry is empty, skip it
+            if not entry:
+                continue
+            if entry["state"] != "missing":
+                tmp_cache[entry["id_path"]] = entry
+                continue
+            if entry["state"] == "missing" and entry["id_path"] not in tmp_cache:
+                tmp_cache[entry["id_path"]] = entry
+    # Convert tmp_cache to list
+    cache = list(tmp_cache.values())
+    cache.sort(key=lambda x: x["id_path"])
+    return cache
+
+
 def get_location_cache(nodes: List[NmNode], id_paths: Dict):
     """
     Get the location cache for a list of NmNode objects.
diff --git a/tests/test_location.py b/tests/test_location.py
new file mode 100644
index 0000000..b5cdbc2
--- /dev/null
+++ b/tests/test_location.py
@@ -0,0 +1,53 @@
+import pytest
+
+import inet_nm.location as loc
+
+
+@pytest.mark.parametrize(
+    "caches",
+    [
+        [[{}]],
+        [[{"id_path": "1", "node_uid": "1", "state": "attached"}]],
+        [
+            [
+                {"id_path": "1", "node_uid": "1", "state": "attached"},
+            ],
+            [{"id_path": "1", "node_uid": "1", "state": "missing"}],
+        ],
+        [
+            [
+                {"id_path": "1", "node_uid": "1", "state": "unassigned"},
+                {"id_path": "2", "node_uid": "2", "state": "missing"},
+            ],
+            [{"id_path": "2", "node_uid": "2", "state": "attached"}],
+        ],
+    ],
+)
+def test_merge_location_cache_chunks_no_missing(caches):
+    cache = loc.merge_location_cache_chunks(caches)
+    assert not any(entry["state"] == "missing" for entry in cache), cache
+
+
+@pytest.mark.parametrize(
+    "caches",
+    [
+        [[{"id_path": "1", "node_uid": "1", "state": "missing"}]],
+        [
+            [
+                {"id_path": "1", "node_uid": "1", "state": "attached"},
+                {"id_path": "2", "node_uid": "2", "state": "missing"},
+            ],
+        ],
+        [
+            [
+                {"id_path": "1", "node_uid": "1", "state": "missing"},
+            ],
+            [
+                {"id_path": "1", "node_uid": "1", "state": "missing"},
+            ],
+        ],
+    ],
+)
+def test_merge_location_cache_chunks_missing(caches):
+    cache = loc.merge_location_cache_chunks(caches)
+    assert any(entry["state"] == "missing" for entry in cache), cache