diff --git a/cli/fix/vars_main_files.py b/cli/fix/vars_main_files.py
index af9a69b0..f4bad3da 100644
--- a/cli/fix/vars_main_files.py
+++ b/cli/fix/vars_main_files.py
@@ -8,6 +8,10 @@ import sys
 import yaml
 from pathlib import Path
 
+# Directory containing roles; can be overridden by tests
+MODULE_DIR = Path(__file__).resolve().parent
+ROLES_DIR = (MODULE_DIR.parent / "roles").resolve()
+
 def process_role(role_dir: Path, prefix: str, preview: bool, overwrite: bool):
     name = role_dir.name
     if not name.startswith(prefix):
@@ -50,6 +54,15 @@ def process_role(role_dir: Path, prefix: str, preview: bool, overwrite: bool):
             print(f"Created {vars_file} with application_id: {expected_id}")
 
 
+def run(prefix: str, preview: bool = False, overwrite: bool = False):
+    """
+    Ensure vars/main.yml for roles under ROLES_DIR with the given prefix has correct application_id.
+    """
+    for role in sorted(Path(ROLES_DIR).iterdir()):
+        if role.is_dir():
+            process_role(role, prefix, preview, overwrite)
+
+
 def main():
     parser = argparse.ArgumentParser(
         description="Ensure vars/main.yml for roles with a given prefix has correct application_id"
@@ -68,16 +81,9 @@ def main():
     )
     args = parser.parse_args()
 
-    # Determine roles directory relative to this script
-    script_dir = Path(__file__).resolve().parent
-    roles_dir = (script_dir.parent / "../roles").resolve()
-    if not roles_dir.is_dir():
-        print(f"Roles directory not found: {roles_dir}", file=sys.stderr)
-        sys.exit(1)
+    # Run processing
+    run(prefix=args.prefix, preview=args.preview, overwrite=args.overwrite)
 
-    for role in sorted(roles_dir.iterdir()):
-        if role.is_dir():
-            process_role(role, args.prefix, args.preview, args.overwrite)
 
 if __name__ == "__main__":
     main()
diff --git a/tests/unit/cli/fix/test_vars_main_files.py b/tests/unit/cli/fix/test_vars_main_files.py
index e7d6f608..30c0f112 100644
--- a/tests/unit/cli/fix/test_vars_main_files.py
+++ b/tests/unit/cli/fix/test_vars_main_files.py
@@ -6,7 +6,7 @@ import unittest
 import yaml
 
 # Adjust this import to match the real path in your project
-from cli.ensure.vars_main_files import run, ROLES_DIR
+from cli.fix.vars_main_files import run, ROLES_DIR
 
 class TestEnsureVarsMain(unittest.TestCase):
     def setUp(self):
@@ -17,11 +17,11 @@ class TestEnsureVarsMain(unittest.TestCase):
 
         # Monkey-patch the module's ROLES_DIR to point here
         self._orig_roles_dir = ROLES_DIR
-        setattr(__import__("cli.ensure.vars_main_files", fromlist=["ROLES_DIR"]), "ROLES_DIR", self.roles_dir)
+        setattr(__import__("cli.fix.vars_main_files", fromlist=["ROLES_DIR"]), "ROLES_DIR", self.roles_dir)
 
     def tearDown(self):
         # restore and cleanup
-        setattr(__import__("cli.ensure.vars_main_files", fromlist=["ROLES_DIR"]), "ROLES_DIR", self._orig_roles_dir)
+        setattr(__import__("cli.fix.vars_main_files", fromlist=["ROLES_DIR"]), "ROLES_DIR", self._orig_roles_dir)
         shutil.rmtree(self.tmpdir)
 
     def _make_role(self, name, vars_content=None):