From b867f962ef8284cb4d1e594eae0740f4da81ec73 Mon Sep 17 00:00:00 2001
From: Max Marrone <max@opentrons.com>
Date: Tue, 21 Jan 2020 14:21:27 -0500
Subject: [PATCH] refactor(update-server): Don't add extra newlines to
 /etc/machine-info on update (#4671)

---
 .../otupdate/buildroot/name_management.py     | 21 +++++++---
 .../tests/buildroot/test_name_management.py   | 42 +++++++++++++++++++
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/update-server/otupdate/buildroot/name_management.py b/update-server/otupdate/buildroot/name_management.py
index 596dc99be67..66c9813c2a3 100644
--- a/update-server/otupdate/buildroot/name_management.py
+++ b/update-server/otupdate/buildroot/name_management.py
@@ -182,15 +182,26 @@ def _update_pretty_hostname(new_val: str):
     except OSError:
         LOG.exception("Couldn't read /etc/machine-info")
         contents = ''
-    new_contents = ''
-    for line in contents.split('\n'):
-        if not line.startswith('PRETTY_HOSTNAME'):
-            new_contents += f'{line}\n'
-    new_contents += f'PRETTY_HOSTNAME={new_val}\n'
+    new_contents = _rewrite_machine_info(contents, new_val)
     with open('/etc/machine-info', 'w') as emi:
         emi.write(new_contents)
 
 
+def _rewrite_machine_info(
+        current_machine_info_contents: str, new_pretty_hostname: str) -> str:
+    """
+    Return current_machine_info_contents - the full contents of
+    /etc/machine-info - with the PRETTY_HOSTNAME=... line rewritten to refer
+    to new_pretty_hostname.
+    """
+    current_lines = current_machine_info_contents.splitlines()
+    preserved_lines = [
+        l for l in current_lines if not l.startswith('PRETTY_HOSTNAME')]
+    new_lines = preserved_lines + [f'PRETTY_HOSTNAME={new_pretty_hostname}']
+    new_contents = '\n'.join(new_lines) + '\n'
+    return new_contents
+
+
 def get_name(default: str = 'no name set'):
     """ Get the currently-configured name of the machine """
     try:
diff --git a/update-server/tests/buildroot/test_name_management.py b/update-server/tests/buildroot/test_name_management.py
index 3b6bb8f7915..dcb107251ed 100644
--- a/update-server/tests/buildroot/test_name_management.py
+++ b/update-server/tests/buildroot/test_name_management.py
@@ -1,4 +1,6 @@
 from otupdate.buildroot import name_management
+import pytest
+from collections import Counter
 
 
 async def test_name_endpoint(test_cli, monkeypatch):
@@ -41,3 +43,43 @@ async def fake_name(name):
     health = await test_cli.get('/server/update/health')
     health_body = await health.json()
     assert health_body['name'] == to_set + to_set
+
+
+machine_info_examples = [
+    '',
+    'FOO=foo',
+    'PRETTY_HOSTNAME=initial_pretty_hostname',
+    'FOO=foo\nPRETTY_HOSTNAME=initial_pretty_hostname\nBAR=bar'
+]
+
+
+@pytest.mark.parametrize('initial_contents', machine_info_examples)
+def test_rewrite_machine_info_updates_pretty_hostname(initial_contents):
+    rewrite = name_management._rewrite_machine_info(
+        initial_contents, 'new_pretty_hostname')
+    assert 'PRETTY_HOSTNAME=new_pretty_hostname' in rewrite.splitlines(), \
+        'new PRETTY_HOSTNAME should be present.'
+    assert rewrite.count('PRETTY_HOSTNAME') == 1, \
+        'Old PRETTY_HOSTNAME should be deleted.'
+
+
+@pytest.mark.parametrize('initial_contents', machine_info_examples)
+def test_rewrite_machine_info_preserves_other_lines(initial_contents):
+    intitial_lines = Counter(initial_contents.splitlines())
+    rewrite_string = name_management._rewrite_machine_info(
+        initial_contents, 'new_pretty_hostname')
+    rewrite_lines = Counter(rewrite_string.splitlines())
+    lost_lines = intitial_lines - rewrite_lines
+    for line in lost_lines:
+        # Lines are only allowed to be "lost" in the rewrite if they were an
+        # old PRETTY_HOSTNAME assignment, or were blank.
+        assert line.startswith('PRETTY_HOSTNAME=') or line == ''
+
+
+@pytest.mark.parametrize('initial_contents', machine_info_examples)
+def test_rewrite_machine_info_is_idempotent(initial_contents):
+    first_rewrite = name_management._rewrite_machine_info(
+        initial_contents, 'new_pretty_hostname')
+    second_rewrite = name_management._rewrite_machine_info(
+        first_rewrite, 'new_pretty_hostname')
+    assert second_rewrite == first_rewrite