Fix cli tests

umsi-mads · Apr 9, 2024 · bbf7f87 · bbf7f87
1 parent 2b88726
commit bbf7f87
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/test/mads/test_cli.py b/test/mads/test_cli.py
@@ -1,17 +1,34 @@
 """Test the CLI interface"""
 
 from mads import cli
+from mads.environ import Git
 
 
 def test_determine_tag(capsys):
     """Test that the help message works"""
 
     cli.main(["docker", "tag"])
-    assert capsys.readouterr().out == "dev\n"
+
+    git = Git()
+    expected = {
+        "main": "latest",
+        "master": "latest",
+        "beta": "beta",
+        None: "dev",
+    }.get(git.branch, "dev")
+    assert capsys.readouterr().out == f"{expected}\n"
 
 
 def test_determine_tag_default(capsys):
     """Test that the help message works"""
 
     cli.main(["docker", "tag", "--default", "beta"])
-    assert capsys.readouterr().out == "beta\n"
+
+    git = Git()
+    expected = {
+        "main": "latest",
+        "master": "latest",
+        "beta": "beta",
+        None: "beta",
+    }.get(git.branch, "beta")
+    assert capsys.readouterr().out == f"{expected}\n"