Fix issues for sonic_installer upgrade-docker and sonic_installer rol…

…lback-docker (#2278)

#### What I did

Fix issues:
1. upgrade-docker shall not check STATE DB for those container which do not support warm mode
2. rollback-docker sometimes cannot really rollback the container because it cannot get the correct docker image tag

#### How I did it

1. for containers do not support warm mode, ignore STATE_DB check
2. use another command to get docker image tag

#### How to verify it

Manual test.

UT is not created for this change because there are many docker command running for these two CLIs, thus:
1. Mock so many docker command causes the test case very tricky. The case would be even more complicated than the CLIs.
2. UT cannot really test the command in case of so many docker command

Instead, a new sonic-mgmt test case is on the way to cover these two CLI. Will create PR for the sonic-mgmt new test case later.

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