From a17dba4c839318b4b36f447d134f4bea28213400 Mon Sep 17 00:00:00 2001
From: Oleg Gelbukh <ogelbukh@mirantis.com>
Date: Sat, 21 Nov 2015 03:12:56 +0000
Subject: [PATCH] Add functions to set/unset upgrade_levels for compute service

These functions set and unset upgrade_levels for the compute service
in nova.conf file. It will be used to handle rolling upgrade of compute
nodes as well as the cleanup of nova.conf once upgrade is finished.

Related-bug: #1506041
Co-Authored-By: Ilya Kharin <akscram@gmail.com>
Change-Id: I4fcefae2051221ad220c0cb70226083060c5aa23
---
 octane/magic_consts.py         |  8 ++++++
 octane/tests/test_util_node.py | 45 ++++++++++++++++++++++++++++++++++
 octane/util/node.py            | 18 ++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/octane/magic_consts.py b/octane/magic_consts.py
index 25c4a544..1c06797b 100644
--- a/octane/magic_consts.py
+++ b/octane/magic_consts.py
@@ -27,3 +27,11 @@ BRIDGES = ['br-ex', 'br-mgmt']
 DEFAULT_DISKS = True
 DEFAULT_NETS = True
 ISCSI_CONFIG_PATH = "/etc/iscsi/initiatorname.iscsi"
+VERSIONS = {
+    '7.0': 'kilo',
+    '6.1': 'juno',
+    '6.0': 'juno',
+    '5.2.9': 'icehouse',
+    '5.1.1': 'icehouse',
+    '5.1': 'icehouse',
+}
diff --git a/octane/tests/test_util_node.py b/octane/tests/test_util_node.py
index 4b685e5a..62517916 100644
--- a/octane/tests/test_util_node.py
+++ b/octane/tests/test_util_node.py
@@ -10,6 +10,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import contextlib
 import io
 import mock
 import pytest
@@ -96,3 +97,47 @@ def test_untar_files(node, mock_ssh_popen, mock_open):
                                            stdin=ssh.PIPE, node=node)
     mock_open.assert_called_once_with('filename', 'rb')
     assert buf.getvalue() == content
+
+
+@contextlib.contextmanager
+def _check_upgrade_levels(mocker, node, content, expected_content):
+    mock_sftp = mocker.patch('octane.util.ssh.sftp')
+
+    old = io.BytesIO(content)
+    mock_new = mocker.Mock()
+    buf = io.BytesIO()
+    mock_new.write = buf.write
+    mock_update_file = mocker.patch('octane.util.ssh.update_file')
+    mock_update_file.return_value.__enter__.return_value = (old, mock_new)
+
+    yield
+
+    mock_sftp.assert_called_once_with(node)
+    mock_update_file.assert_called_once_with(mock_sftp.return_value,
+                                             '/etc/nova/nova.conf')
+    assert buf.getvalue() == expected_content
+
+
+NOVA_DEFAULT = b"#\u0444\n[DEFAULT]\ndebug = True\n"
+NOVA_WITH_EMPTY_LEVELS = NOVA_DEFAULT + b"[upgrade_levels]\n"
+NOVA_WITH_KILO_LEVELS = NOVA_WITH_EMPTY_LEVELS + b"compute=kilo\n"
+
+
+@pytest.mark.parametrize("content,expected_content", [
+    (NOVA_DEFAULT, NOVA_DEFAULT),
+    (NOVA_WITH_EMPTY_LEVELS, NOVA_WITH_KILO_LEVELS),
+])
+def test_add_compute_upgrade_levels(mocker, node, content, expected_content):
+    with _check_upgrade_levels(mocker, node, content, expected_content):
+        node_util.add_compute_upgrade_levels(node, 'kilo')
+
+
+@pytest.mark.parametrize("content,expected_content", [
+    (NOVA_DEFAULT, NOVA_DEFAULT),
+    (NOVA_WITH_EMPTY_LEVELS, NOVA_WITH_EMPTY_LEVELS),
+    (NOVA_WITH_KILO_LEVELS, NOVA_WITH_EMPTY_LEVELS),
+])
+def test_remove_compute_upgrade_levels(mocker, node, content,
+                                       expected_content):
+    with _check_upgrade_levels(mocker, node, content, expected_content):
+        node_util.remove_compute_upgrade_levels(node)
diff --git a/octane/util/node.py b/octane/util/node.py
index 78cbdcc5..359d01bf 100644
--- a/octane/util/node.py
+++ b/octane/util/node.py
@@ -119,3 +119,21 @@ def wait_for_mcollective_start(nodes, timeout=600):
             raise Exception("Timeout waiting for nodes {0} to start"
                             " mcollective".format(failed))
         wait_list -= done
+
+
+def add_compute_upgrade_levels(node, version):
+    sftp = ssh.sftp(node)
+    with ssh.update_file(sftp, '/etc/nova/nova.conf') as (old, new):
+        for line in old:
+            new.write(line)
+            if line.startswith("[upgrade_levels]"):
+                new.write("compute={0}\n".format(version))
+
+
+def remove_compute_upgrade_levels(node):
+    sftp = ssh.sftp(node)
+    with ssh.update_file(sftp, '/etc/nova/nova.conf') as (old, new):
+        for line in old:
+            if line.startswith("compute="):
+                continue
+            new.write(line)