Fix review comment and add unit test

sonic_installer/common.py

@@ -31,8 +31,6 @@ def run_command(command):
     if proc.returncode != 0:
         sys.exit(proc.returncode)
 
-    return out.rstrip("\n")
-
 # Run bash command and return output, raise if it fails
 def run_command_or_raise(argv, raise_exception=True):
     click.echo(click.style("Command: ", fg='cyan') + click.style(' '.join(argv), fg='green'))

sonic_installer/main.py

@@ -787,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:
@@ -867,16 +859,12 @@ 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)
-
     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_latest and id == image_id_previous:
+            if id == image_id_previous:
                 run_command("docker rmi -f %s" % id)
                 break
 
@@ -939,7 +927,7 @@ def rollback_docker(container_name):
     version_tag = ""
     for id in image_id_all:
         if id != image_id_previous:
-            version_tag = run_command("docker images --format '{{{{.ID}}}} {{{{.Tag}}}}' | grep {} | awk '{{print $2}}'".format(id))
+            version_tag = get_docker_tag_name(id)
             break
 
     # make previous image as latest

tests/installer_docker_test.py

@@ -0,0 +1,126 @@
+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.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