diff --git a/src/_nebari/upgrade.py b/src/_nebari/upgrade.py index 6cb5b098a6..b0d5b472b3 100644 --- a/src/_nebari/upgrade.py +++ b/src/_nebari/upgrade.py @@ -171,7 +171,7 @@ def upgrade_step(self, config, start_version, config_filename, *args, **kwargs): def contains_image_and_tag(s: str) -> bool: # match on `quay.io/nebari/nebari-<...>:YYYY.MM.XX`` - pattern = r"^quay\.io\/nebari\/nebari-(jupyterhub|jupyterlab|dask-worker):\d{4}\.\d+\.\d+$" + pattern = r"^quay\.io\/nebari\/nebari-(jupyterhub|jupyterlab|dask-worker)(-gpu)?:\d{4}\.\d+\.\d+$" return bool(re.match(pattern, s)) def replace_image_tag_legacy(image, start_version, new_version): @@ -249,11 +249,11 @@ def update_image_tag(config, config_path, current_image, new_version): # update profiles.dask_worker images for k, v in config.get("profiles", {}).get("dask_worker", {}).items(): - current_image = v.get("kubespawner_override", {}).get("image", None) + current_image = v.get("image", None) if current_image: config = update_image_tag( config, - f"profiles.dask_worker.{k}.kubespawner_override.image", + f"profiles.dask_worker.{k}.image", current_image, __rounded_finish_version__, ) diff --git a/tests/tests_unit/test_cli_upgrade.py b/tests/tests_unit/test_cli_upgrade.py new file mode 100644 index 0000000000..d00c04b1b0 --- /dev/null +++ b/tests/tests_unit/test_cli_upgrade.py @@ -0,0 +1,363 @@ +import tempfile +from pathlib import Path +from typing import Any, Dict, List + +import pytest +import yaml +from typer.testing import CliRunner + +import _nebari.upgrade +import _nebari.version +from _nebari.cli import create_cli + + +# can't upgrade to a previous version that doesn't have a corresponding +# UpgradeStep derived class. without these dummy classes, the rendered +# nebari-config.yaml will have the wrong version +class Test_Cli_Upgrade_2022_10_1(_nebari.upgrade.UpgradeStep): + version = "2022.10.1" + + +class Test_Cli_Upgrade_2022_11_1(_nebari.upgrade.UpgradeStep): + version = "2022.11.1" + + +class Test_Cli_Upgrade_2023_1_1(_nebari.upgrade.UpgradeStep): + version = "2023.1.1" + + +class Test_Cli_Upgrade_2023_4_1(_nebari.upgrade.UpgradeStep): + version = "2023.4.1" + + +class Test_Cli_Upgrade_2023_5_1(_nebari.upgrade.UpgradeStep): + version = "2023.5.1" + + +### end dummy upgrade classes + +runner = CliRunner() + + +@pytest.mark.parametrize( + "args, exit_code, content", + [ + # --help + (["--help"], 0, ["Usage:"]), + (["-h"], 0, ["Usage:"]), + # error, missing args + ([], 2, ["Missing option"]), + (["--config"], 2, ["requires an argument"]), + (["-c"], 2, ["requires an argument"]), + # only used for old qhub version upgrades + ( + ["--attempt-fixes"], + 2, + ["Missing option"], + ), + ], +) +def test_upgrade_stdout(args: List[str], exit_code: int, content: List[str]): + app = create_cli() + result = runner.invoke(app, ["upgrade"] + args) + assert result.exit_code == exit_code + for c in content: + assert c in result.stdout + + +def test_upgrade_2022_10_1_to_2022_11_1(monkeypatch: pytest.MonkeyPatch): + assert_nebari_upgrade_success(monkeypatch, "2022.10.1", "2022.11.1") + + +def test_upgrade_2022_11_1_to_2023_1_1(monkeypatch: pytest.MonkeyPatch): + assert_nebari_upgrade_success(monkeypatch, "2022.11.1", "2023.1.1") + + +def test_upgrade_2023_1_1_to_2023_4_1(monkeypatch: pytest.MonkeyPatch): + assert_nebari_upgrade_success(monkeypatch, "2023.1.1", "2023.4.1") + + +def test_upgrade_2023_4_1_to_2023_5_1(monkeypatch: pytest.MonkeyPatch): + assert_nebari_upgrade_success( + monkeypatch, + "2023.4.1", + "2023.5.1", + # Have you deleted the Argo Workflows CRDs and service accounts + inputs=["y"], + ) + + +@pytest.mark.parametrize( + "workflows_enabled, workflow_controller_enabled", + [(True, True), (True, False), (False, None), (None, None)], +) +def test_upgrade_2023_5_1_to_2023_7_2( + monkeypatch: pytest.MonkeyPatch, + workflows_enabled: bool, + workflow_controller_enabled: bool, +): + addl_config = {} + inputs = [] + + if workflows_enabled is not None: + addl_config = {"argo_workflows": {"enabled": workflows_enabled}} + if workflows_enabled is True: + inputs.append("y" if workflow_controller_enabled else "n") + + upgraded = assert_nebari_upgrade_success( + monkeypatch, + "2023.5.1", + "2023.7.2", + addl_config=addl_config, + # Do you want to enable the Nebari Workflow Controller? + inputs=inputs, + ) + + if workflows_enabled is True: + if workflow_controller_enabled: + assert ( + True + is upgraded["argo_workflows"]["nebari_workflow_controller"]["enabled"] + ) + else: + # not sure this makes sense, the code defaults this to True if missing + assert "nebari_workflow_controller" not in upgraded["argo_workflows"] + elif workflows_enabled is False: + assert False is upgraded["argo_workflows"]["enabled"] + else: + # argo_workflows config missing + # this one doesn't sound right either, they default to true if this is missing, so why skip the questions? + assert "argo_workflows" not in upgraded + + +def test_upgrade_image_tags(monkeypatch: pytest.MonkeyPatch): + start_version = "2023.5.1" + end_version = "2023.7.2" + + upgraded = assert_nebari_upgrade_success( + monkeypatch, + start_version, + end_version, + # # number of "y" inputs directly corresponds to how many matching images are found in yaml + inputs=["y", "y", "y", "y", "y", "y", "y"], + addl_config=yaml.safe_load( + f""" +default_images: + jupyterhub: quay.io/nebari/nebari-jupyterhub:{start_version} + jupyterlab: quay.io/nebari/nebari-jupyterlab:{start_version} + dask_worker: quay.io/nebari/nebari-dask-worker:{start_version} +profiles: + jupyterlab: + - display_name: base + kubespawner_override: + image: quay.io/nebari/nebari-jupyterlab:{start_version} + - display_name: gpu + kubespawner_override: + image: quay.io/nebari/nebari-jupyterlab-gpu:{start_version} + - display_name: any-other-version + kubespawner_override: + image: quay.io/nebari/nebari-jupyterlab:1955.11.5 + - display_name: leave-me-alone + kubespawner_override: + image: quay.io/nebari/leave-me-alone:{start_version} + dask_worker: + test: + image: quay.io/nebari/nebari-dask-worker:{start_version} +""" + ), + ) + + for _, v in upgraded["default_images"].items(): + assert v.endswith(end_version) + + for profile in upgraded["profiles"]["jupyterlab"]: + if profile["display_name"] != "leave-me-alone": + # assume all other images should have been upgraded to the end_version + assert profile["kubespawner_override"]["image"].endswith(end_version) + else: + # this one was selected not to match the regex for nebari images, should have been left alone + assert profile["kubespawner_override"]["image"].endswith(start_version) + + for _, profile in upgraded["profiles"]["dask_worker"].items(): + assert profile["image"].endswith(end_version) + + +def test_upgrade_fail_on_missing_file(): + with tempfile.TemporaryDirectory() as tmp: + tmp_file = Path(tmp).resolve() / "nebari-config.yaml" + assert tmp_file.exists() is False + + app = create_cli() + + result = runner.invoke(app, ["upgrade", "--config", tmp_file.resolve()]) + + assert 1 == result.exit_code + assert result.exception + assert ( + f"passed in configuration filename={tmp_file.resolve()} must exist" + in str(result.exception) + ) + + +def test_upgrade_fail_invalid_file(): + with tempfile.TemporaryDirectory() as tmp: + tmp_file = Path(tmp).resolve() / "nebari-config.yaml" + assert tmp_file.exists() is False + + nebari_config = yaml.safe_load( + """ +project_name: test +provider: fake + """ + ) + + with open(tmp_file.resolve(), "w") as f: + yaml.dump(nebari_config, f) + + assert tmp_file.exists() is True + app = create_cli() + + result = runner.invoke(app, ["upgrade", "--config", tmp_file.resolve()]) + + assert 1 == result.exit_code + assert result.exception + assert "provider" in str(result.exception) + + +def test_upgrade_fail_on_downgrade(): + start_version = "9999.9.9" # way in the future + end_version = _nebari.upgrade.__version__ + + with tempfile.TemporaryDirectory() as tmp: + tmp_file = Path(tmp).resolve() / "nebari-config.yaml" + assert tmp_file.exists() is False + + nebari_config = yaml.safe_load( + f""" +project_name: test +provider: local +domain: test.example.com +namespace: dev +nebari_version: {start_version} + """ + ) + + with open(tmp_file.resolve(), "w") as f: + yaml.dump(nebari_config, f) + + assert tmp_file.exists() is True + app = create_cli() + + result = runner.invoke(app, ["upgrade", "--config", tmp_file.resolve()]) + + assert 1 == result.exit_code + assert result.exception + assert ( + f"already belongs to a later version ({start_version}) than the installed version of Nebari ({end_version})" + in str(result.exception) + ) + + # make sure the file is unaltered + with open(tmp_file.resolve(), "r") as c: + assert yaml.safe_load(c) == nebari_config + + +def test_upgrade_does_nothing_on_same_version(): + # this test only seems to work against the actual current version, any + # mocked earlier versions trigger an actual update + start_version = _nebari.upgrade.__version__ + + with tempfile.TemporaryDirectory() as tmp: + tmp_file = Path(tmp).resolve() / "nebari-config.yaml" + assert tmp_file.exists() is False + + nebari_config = yaml.safe_load( + f""" +project_name: test +provider: local +domain: test.example.com +namespace: dev +nebari_version: {start_version} + """ + ) + + with open(tmp_file.resolve(), "w") as f: + yaml.dump(nebari_config, f) + + assert tmp_file.exists() is True + app = create_cli() + + result = runner.invoke(app, ["upgrade", "--config", tmp_file.resolve()]) + + # feels like this should return a non-zero exit code if the upgrade is not happening + assert 0 == result.exit_code + assert not result.exception + assert "up-to-date" in result.stdout + + # make sure the file is unaltered + with open(tmp_file.resolve(), "r") as c: + assert yaml.safe_load(c) == nebari_config + + +def assert_nebari_upgrade_success( + monkeypatch: pytest.MonkeyPatch, + start_version: str, + end_version: str, + addl_config: Dict[str, Any] = {}, + inputs: List[str] = [], +) -> Dict[str, Any]: + monkeypatch.setattr(_nebari.upgrade, "__version__", end_version) + + # create a tmp dir and clean up when done + with tempfile.TemporaryDirectory() as tmp: + tmp_file = Path(tmp).resolve() / "nebari-config.yaml" + assert tmp_file.exists() is False + + # merge basic config with any test case specific values provided + nebari_config = { + **yaml.safe_load( + f""" +project_name: test +provider: local +domain: test.example.com +namespace: dev +nebari_version: {start_version} + """ + ), + **addl_config, + } + + # write the test nebari-config.yaml file to tmp location + with open(tmp_file.resolve(), "w") as f: + yaml.dump(nebari_config, f) + + assert tmp_file.exists() is True + app = create_cli() + + if inputs is not None and len(inputs) > 0: + inputs.append("") # trailing newline for last input + + # run nebari upgrade -c tmp/nebari-config.yaml + result = runner.invoke( + app, ["upgrade", "--config", tmp_file.resolve()], input="\n".join(inputs) + ) + + assert 0 == result.exit_code + assert not result.exception + assert "Saving new config file" in result.stdout + + # load the modified nebari-config.yaml and check the new version has changed + with open(tmp_file.resolve(), "r") as f: + upgraded = yaml.safe_load(f) + assert end_version == upgraded["nebari_version"] + + # check backup matches original + backup_file = Path(tmp).resolve() / f"nebari-config.yaml.{start_version}.backup" + assert backup_file.exists() is True + with open(backup_file.resolve(), "r") as b: + backup = yaml.safe_load(b) + assert backup == nebari_config + + # pass the parsed nebari-config.yaml with upgrade mods back to caller for + # additional assertions + return upgraded