From 03c40528d4300e465a48d3dbaabe1c04fa5734dc Mon Sep 17 00:00:00 2001
From: Ian Fijolek <ian@iamthefij.com>
Date: Thu, 3 Oct 2024 15:47:21 -0700
Subject: [PATCH] Add freeze flag to list command to regenerate unhacs.yaml
 file

---
 tests/main_test.py | 25 +++++++++++++++++++++++++
 unhacs/main.py     | 16 +++++++++++++---
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/tests/main_test.py b/tests/main_test.py
index 887b8c2..c369fba 100644
--- a/tests/main_test.py
+++ b/tests/main_test.py
@@ -6,6 +6,7 @@ from pathlib import Path
 
 from unhacs.main import main
 from unhacs.packages import get_installed_packages
+from unhacs.packages import read_lock_packages
 
 INTEGRATION_URL = "https://github.com/simbaja/ha_gehome"
 INTEGRATION_VERSION = "v0.6.9"
@@ -124,6 +125,30 @@ class TestMainIntegrarion(unittest.TestCase):
         self.assertEqual(installed[0].url, INTEGRATION_URL)
         self.assertEqual(installed[0].version, INTEGRATION_VERSION)
 
+        # Delete the custom_components folder and re-install the integration using the lock file
+        shutil.rmtree(os.path.join(self.test_dir, "custom_components"))
+        self.run_itest(
+            "Re-install integration using lock file",
+            "add --file unhacs.yaml",
+            expected_files=[
+                "custom_components/ge_home/__init__.py",
+                "custom_components/ge_home/manifest.json",
+                "custom_components/ge_home/switch.py",
+            ],
+        )
+
+        # Delete the lock file and then regenerate it
+        os.remove(os.path.join(self.test_dir, "unhacs.yaml"))
+        self.run_itest(
+            "Regenerate lock file",
+            "list --freeze",
+            expected_files=[
+                "unhacs.yaml",
+            ],
+        )
+
+        self.assertGreater(len(read_lock_packages()), 0)
+
         self.run_itest(
             "Remove integration",
             "remove ha_gehome --yes",
diff --git a/unhacs/main.py b/unhacs/main.py
index fe321b4..9f66d88 100644
--- a/unhacs/main.py
+++ b/unhacs/main.py
@@ -55,6 +55,12 @@ def parse_args(argv: list[str]):
     # List installed packages
     list_parser = subparsers.add_parser("list", description="List installed packages.")
     list_parser.add_argument("--verbose", "-v", action="store_true")
+    list_parser.add_argument(
+        "--freeze",
+        "-f",
+        action="store_true",
+        description="Regenerate unhacs.yaml with installed packages.",
+    )
 
     # List git tags for a given package
     list_tags_parser = subparsers.add_parser("tags", help="List tags for a package.")
@@ -233,11 +239,15 @@ class Unhacs:
 
         self.write_lock_packages(packages)
 
-    def list_packages(self, verbose: bool = False):
+    def list_packages(self, verbose: bool = False, freeze: bool = False):
         """List installed packages and their versions."""
-        for package in get_installed_packages():
+        installed_packages = get_installed_packages()
+        for package in installed_packages:
             print(package.verbose_str() if verbose else str(package))
 
+        if freeze:
+            self.write_lock_packages(installed_packages)
+
     def list_tags(self, url: str, limit: int = 10):
         print(f"Tags for {url}:")
         for tag in get_repo_tags(url)[-1 * limit :]:
@@ -343,7 +353,7 @@ def main(argv: list[str] | None = None) -> int:
             print("Either a file or a URL must be provided")
             return 1
     elif args.subcommand == "list":
-        unhacs.list_packages(args.verbose)
+        unhacs.list_packages(args.verbose, args.freeze)
     elif args.subcommand == "tags":
         unhacs.list_tags(args.url, limit=args.limit)
     elif args.subcommand == "remove":