feat: allow setting different root dir for apt

Signed-off-by: egvimo <egvimo@gmail.com>

Author: egvimo

apt_info.py

@@ -8,10 +8,11 @@
 # Authors: Kyle Fazzari <kyrofa@ubuntu.com>
 #          Daniel Swarbrick <dswarbrick@debian.org>
 
-import apt
+import argparse
 import collections
 import contextlib
 import os
+import apt
 from prometheus_client import CollectorRegistry, Gauge, generate_latest
 
 _UpgradeInfo = collections.namedtuple("_UpgradeInfo", ["labels", "count"])
@@ -26,12 +27,12 @@ def _convert_candidates_to_upgrade_infos(candidates):
         )
         changes_dict[",".join(origins)][candidate.architecture] += 1
 
-    changes_list = list()
+    changes_list = []
     for origin in sorted(changes_dict.keys()):
         for arch in sorted(changes_dict[origin].keys()):
             changes_list.append(
                 _UpgradeInfo(
-                    labels=dict(origin=origin, arch=arch),
+                    labels={"origin": origin, "arch": arch},
                     count=changes_dict[origin][arch],
                 )
             )
@@ -74,14 +75,14 @@ def _write_autoremove_pending(registry, cache):
     g.set(len(autoremovable_packages))
 
 
-def _write_reboot_required(registry):
+def _write_reboot_required(registry, root_dir):
     g = Gauge('node_reboot_required', "Node reboot is required for software updates.",
               registry=registry)
-    g.set(int(os.path.isfile('/run/reboot-required')))
+    g.set(int(os.path.isfile(os.path.join(root_dir, 'run/reboot-required'))))
 
 
-def _main():
-    cache = apt.cache.Cache()
+def generate_metrics(root_dir: str = '/') -> bytes:
+    cache = apt.cache.Cache(rootdir=root_dir)
 
     # First of all, attempt to update the index. If we don't have permission
     # to do so (or it fails for some reason), it's not the end of the world,
@@ -96,8 +97,20 @@ def _main():
     _write_pending_upgrades(registry, cache)
     _write_held_upgrades(registry, cache)
     _write_autoremove_pending(registry, cache)
-    _write_reboot_required(registry)
-    print(generate_latest(registry).decode(), end='')
+    _write_reboot_required(registry, root_dir)
+
+    return generate_latest(registry)
+
+
+def _main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-r', '--root-dir', dest='root_dir', type=str, default='/',
+                        help="Set root directory to a different path than /")
+    args = parser.parse_args()
+
+    metrics = generate_metrics(args.root_dir)
+
+    print(metrics.decode(), end='')
 
 
 if __name__ == "__main__":