Fix issues for sonic_installer upgrade-docker and sonic_installer rollback-docker

sonic_installer/main.py

@@ -134,7 +134,7 @@ def get_docker_tag_name(image):
 
 
 def echo_and_log(msg, priority=LOG_NOTICE, fg=None):
-    if priority >= LOG_ERR:
+    if priority == LOG_ERR:
         # Print to stderr if priority is error
         click.secho(msg, fg=fg, err=True)
     else:
@@ -647,7 +647,7 @@ def set_fips(image, enable_fips):
     bootloader = get_bootloader()
     if not image:
         image =  bootloader.get_next_image()
-    if image not in bootloader.get_installed_images(): 
+    if image not in bootloader.get_installed_images():
         echo_and_log('Error: Image does not exist', LOG_ERR)
         sys.exit(1)
     bootloader.set_fips(image, enable=enable_fips)
@@ -743,7 +743,8 @@ def cleanup():
     "swss",
     "syncd",
     "teamd",
-    "telemetry"
+    "telemetry",
+    "mgmt-framework"
 ]
 
 # Upgrade docker image
@@ -786,16 +787,8 @@ def upgrade_docker(container_name, url, cleanup_image, skip_check, tag, warm):
         echo_and_log("Image file '{}' does not exist or is not a regular file. Aborting...".format(image_path), LOG_ERR)
         raise click.Abort()
 
-    warm_configured = False
     # warm restart enable/disable config is put in stateDB, not persistent across cold reboot, not saved to config_DB.json file
-    state_db = SonicV2Connector(host='127.0.0.1')
-    state_db.connect(state_db.STATE_DB, False)
-    TABLE_NAME_SEPARATOR = '|'
-    prefix = 'WARM_RESTART_ENABLE_TABLE' + TABLE_NAME_SEPARATOR
-    _hash = '{}{}'.format(prefix, container_name)
-    if state_db.get(state_db.STATE_DB, _hash, "enable") == "true":
-        warm_configured = True
-    state_db.close(state_db.STATE_DB)
+    warm_configured = hget_warm_restart_table('STATE_DB', 'WARM_RESTART_ENABLE_TABLE', container_name, 'enable') == "true"
 
     if container_name == "swss" or container_name == "bgp" or container_name == "teamd":
         if warm_configured is False and warm:
@@ -866,23 +859,19 @@ def upgrade_docker(container_name, url, cleanup_image, skip_check, tag, warm):
     run_command("docker tag %s:latest %s:%s" % (image_name, image_name, tag))
     run_command("systemctl restart %s" % container_name)
 
-    # All images id under the image name
-    image_id_all = get_container_image_id_all(image_name)
-
-    # this is image_id for image with "latest" tag
-    image_id_latest = get_container_image_id(image_latest)
-
-    for id in image_id_all:
-        if id != image_id_latest:
-            # Unless requested, the previoud docker image will be preserved
-            if not cleanup_image and id == image_id_previous:
-                continue
-            run_command("docker rmi -f %s" % id)
+    if cleanup_image:
+        # All images id under the image name
+        image_id_all = get_container_image_id_all(image_name)
+        # Unless requested, the previoud docker image will be preserved
+        for id in image_id_all:
+            if id == image_id_previous:
+                run_command("docker rmi -f %s" % id)
+                break
 
     exp_state = "reconciled"
     state = ""
     # post warm restart specific procssing for swss, bgp and teamd dockers, wait for reconciliation state.
-    if warm_configured is True or warm:
+    if warm_app_names and (warm_configured is True or warm):
         count = 0
         for warm_app_name in warm_app_names:
             state = ""
@@ -939,6 +928,7 @@ def rollback_docker(container_name):
     for id in image_id_all:
         if id != image_id_previous:
             version_tag = get_docker_tag_name(id)
+            break
 
     # make previous image as latest
     run_command("docker tag %s:%s %s:latest" % (image_name, version_tag, image_name))

tests/installer_docker_test.py

@@ -0,0 +1,127 @@
+import pytest
+import sonic_installer.main as sonic_installer
+
+from click.testing import CliRunner
+from unittest.mock import patch, MagicMock
+
+SUCCESS = 0
+
+
+@patch('sonic_installer.main.get_container_image_name', MagicMock(return_value='docker-fpm-frr'))
+@patch('sonic_installer.main.get_container_image_id_all', MagicMock(return_value=['1', '2']))
+@patch('sonic_installer.main.get_container_image_id', MagicMock(return_value=['1']))
+@patch('sonic_installer.main.get_docker_tag_name', MagicMock(return_value='some_tag'))
+@patch('sonic_installer.main.echo_and_log', MagicMock())
+@patch('sonic_installer.main.run_command')
+def test_rollback_docker_basic(mock_run_cmd):
+    runner = CliRunner()
+    result = runner.invoke(
+            sonic_installer.sonic_installer.commands['rollback-docker'], ['-y', 'bgp']
+        )
+
+    assert result.exit_code == SUCCESS
+    expect_docker_tag_command = 'docker tag docker-fpm-frr:some_tag docker-fpm-frr:latest'
+    mock_run_cmd.assert_called_with(expect_docker_tag_command)
+
+    mock_run_cmd.reset()
+    result = runner.invoke(
+            sonic_installer.sonic_installer.commands['rollback-docker'], ['-y', 'snmp']
+        )
+
+    assert result.exit_code == SUCCESS
+    mock_run_cmd.assert_any_call('systemctl restart snmp')
+
+
+@patch('sonic_installer.main.get_container_image_name', MagicMock(return_value='docker-fpm-frr'))
+@patch('sonic_installer.main.get_container_image_id_all', MagicMock(return_value=['1']))
+def test_rollback_docker_no_extra_image():
+    runner = CliRunner()
+    result = runner.invoke(
+            sonic_installer.sonic_installer.commands['rollback-docker'], ['-y', 'bgp']
+        )
+    assert result.exit_code != SUCCESS
+
+
+@pytest.mark.parametrize("container", ['bgp', 'swss', 'teamd', 'pmon'])
+@patch('sonic_installer.main.get_container_image_name', MagicMock(return_value='docker-fpm-frr'))
+@patch('sonic_installer.main.get_container_image_id', MagicMock(return_value='1'))
+@patch('sonic_installer.main.get_container_image_id_all', MagicMock(return_value=['1', '2']))
+@patch('sonic_installer.main.validate_url_or_abort', MagicMock())
+@patch('sonic_installer.main.urlretrieve', MagicMock())
+@patch('os.path.isfile', MagicMock(return_value=True))
+@patch('sonic_installer.main.get_docker_tag_name', MagicMock(return_value='some_tag'))
+@patch('sonic_installer.main.run_command', MagicMock())
+@patch("sonic_installer.main.subprocess.Popen")
+@patch('sonic_installer.main.hget_warm_restart_table')
+def test_upgrade_docker_basic(mock_hget, mock_popen, container):
+    def mock_hget_impl(db_name, table_name, warm_app_name, key):
+        if table_name == 'WARM_RESTART_ENABLE_TABLE':
+            return "false"
+        elif table_name == 'WARM_RESTART_TABLE':
+            return 'reconciled'
+
+    mock_hget.side_effect = mock_hget_impl
+    mock_proc = MagicMock()
+    mock_proc.communicate = MagicMock(return_value=(None, None))
+    mock_proc.returncode = 0
+    mock_popen.return_value = mock_proc
+
+    runner = CliRunner()
+    result = runner.invoke(
+            sonic_installer.sonic_installer.commands['upgrade-docker'],
+            ['-y', '--cleanup_image', '--warm', container, 'http://']
+        )
+
+    print(result.output)
+    assert result.exit_code == SUCCESS
+
+
+@patch('sonic_installer.main.get_container_image_name', MagicMock(return_value='docker-fpm-frr'))
+@patch('sonic_installer.main.get_container_image_id', MagicMock(return_value=['1']))
+@patch('sonic_installer.main.validate_url_or_abort', MagicMock())
+@patch('sonic_installer.main.urlretrieve', MagicMock(side_effect=Exception('download failed')))
+def test_upgrade_docker_download_fail():
+    runner = CliRunner()
+    result = runner.invoke(
+            sonic_installer.sonic_installer.commands['upgrade-docker'],
+            ['-y', '--cleanup_image', '--warm', 'bgp', 'http://']
+        )
+    assert 'download failed' in result.output
+    assert result.exit_code != SUCCESS
+
+
+@patch('sonic_installer.main.get_container_image_name', MagicMock(return_value='docker-fpm-frr'))
+@patch('sonic_installer.main.get_container_image_id', MagicMock(return_value=['1']))
+@patch('sonic_installer.main.validate_url_or_abort', MagicMock())
+@patch('sonic_installer.main.urlretrieve', MagicMock(side_effect=Exception('download failed')))
+def test_upgrade_docker_image_not_exist():
+    runner = CliRunner()
+    result = runner.invoke(
+            sonic_installer.sonic_installer.commands['upgrade-docker'],
+            ['-y', '--cleanup_image', '--warm', 'bgp', 'invalid_url']
+        )
+    assert 'does not exist' in result.output
+    assert result.exit_code != SUCCESS
+
+
+@patch('sonic_installer.main.get_container_image_name', MagicMock(return_value='docker-fpm-frr'))
+@patch('sonic_installer.main.get_container_image_id', MagicMock(return_value=['1']))
+@patch('sonic_installer.main.validate_url_or_abort', MagicMock())
+@patch('sonic_installer.main.urlretrieve', MagicMock())
+@patch('os.path.isfile', MagicMock(return_value=True))
+@patch('sonic_installer.main.get_docker_tag_name', MagicMock(return_value='some_tag'))
+@patch('sonic_installer.main.run_command', MagicMock())
+@patch('sonic_installer.main.hget_warm_restart_table', MagicMock(return_value='false'))
+@patch("sonic_installer.main.subprocess.Popen")
+def test_upgrade_docker_image_swss_check_failed(mock_popen):
+    mock_proc = MagicMock()
+    mock_proc.communicate = MagicMock(return_value=(None, None))
+    mock_proc.returncode = 1
+    mock_popen.return_value = mock_proc
+    runner = CliRunner()
+    result = runner.invoke(
+            sonic_installer.sonic_installer.commands['upgrade-docker'],
+            ['-y', '--cleanup_image', '--warm', 'swss', 'http://']
+        )
+    assert 'RESTARTCHECK failed' in result.output
+    assert result.exit_code != SUCCESS