Explicitly pass cli argv to main

ecmwf-lab · Feb 15, 2024 · 922b793 · 922b793
1 parent a67984d
commit 922b793
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/ai_models/__main__.py b/ai_models/__main__.py
@@ -18,7 +18,7 @@
 LOG = logging.getLogger(__name__)
 
 
-def _main():
+def _main(argv):
     parser = argparse.ArgumentParser()
 
     # See https://github.com/pytorch/pytorch/issues/77764
@@ -212,15 +212,15 @@ def _main():
         help="Model version",
     )
 
-    if "--models" not in sys.argv:
+    if "--models" not in argv:
         parser.add_argument(
             "model",
             metavar="MODEL",
             choices=available_models(),
             help="The model to run",
         )
 
-    args, unknownargs = parser.parse_known_args()
+    args, unknownargs = parser.parse_known_args(argv)
 
     if args.models:
         for p in sorted(available_models()):
@@ -302,7 +302,7 @@ def _main():
 
 def main():
     with Timer("Total time"):
-        _main()
+        _main(sys.argv[1:])
 
 
 if __name__ == "__main__":