refactored hosts file update. added unit tests

Makefile

@@ -5,7 +5,7 @@ HOOKS_DIR := $(PWD)/hooks
 TEST_PREFIX := PYTHONPATH=$(HOOKS_DIR)
 
 lint:
-	@flake8 --exclude hooks/charmhelpers hooks
+	@flake8 --exclude hooks/charmhelpers hooks unit_tests
 	@charm proof
 
 bin/charm_helpers_sync.py:
@@ -20,7 +20,7 @@ publish: lint
 	bzr push lp:charms/rabbitmq-server
 	bzr push lp:charms/trusty/rabbitmq-server
 
-test:
+unit_test:
 	@echo Starting tests...
-	CHARM_DIR=$(CHARM_DIR) $(TEST_PREFIX) nosetests $(CHARM_DIR)/hooks
+	CHARM_DIR=$(CHARM_DIR) $(TEST_PREFIX) nosetests unit_tests
 

hooks/rabbit_utils.py

@@ -6,6 +6,7 @@ import sys
 import subprocess
 import glob
 from lib.utils import render_template
+import tempfile
 
 from charmhelpers.contrib.openstack.utils import (
     get_hostname,
@@ -17,6 +18,7 @@ from charmhelpers.core.hookenv import (
     relation_get,
     related_units,
     log, ERROR,
+    INFO,
     service_name
 )
 
@@ -41,6 +43,7 @@ ENV_CONF = '/etc/rabbitmq/rabbitmq-env.conf'
 RABBITMQ_CONF = '/etc/rabbitmq/rabbitmq.config'
 RABBIT_USER = 'rabbitmq'
 LIB_PATH = '/var/lib/rabbitmq/'
+HOSTS_FILE = '/etc/hosts'
 
 _named_passwd = '/var/lib/charm/{}/{}.passwd'
 
@@ -372,21 +375,42 @@ def bind_ipv6_interface():
         conf.write(''.join(out))
 
 
-def render_hosts(ip, hostname):
-    if not ip or not hostname:
-        return
-    FILE = '/etc/hosts'
-    with open(FILE, 'r') as hosts:
+def update_hosts_file(map):
+    """Rabbitmq does not currently like ipv6 addresses so we need to use dns
+    names instead. In order to make them resolvable we ensure they are  in
+    /etc/hosts.
+
+    """
+    with open(HOSTS_FILE, 'r') as hosts:
         lines = hosts.readlines()
 
-    for line in lines:
-        if line.startswith(ip) or hostname in line:
-            lines.remove(line)
-    lines.append(ip + ' ' + hostname + '\n')
+    log("Updating hosts file with: %s (current: %s)" % (map, lines),
+        level=INFO)
 
-    with open(FILE, 'w') as hosts:
+    newlines = []
+    for ip, hostname in map.items():
+        if not ip or not hostname:
+            continue
+
+        keepers = []
         for line in lines:
-            hosts.write(line)
+            _line = line.split()
+            if len(line) < 2 or not (_line[0] == ip or hostname in _line[1:]):
+                keepers.append(line)
+            else:
+                log("Removing line '%s' from hosts file" % (line))
+
+        lines = keepers
+        newlines.append("%s %s\n" % (ip, hostname))
+
+    lines += newlines
+
+    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
+        with open(tmpfile.name, 'w') as hosts:
+            for line in lines:
+                hosts.write(line)
+
+    os.rename(tmpfile.name, HOSTS_FILE)
 
 
 def assert_charm_supports_ipv6():

hooks/rabbitmq_server_relations.py

@@ -215,10 +215,11 @@ def cluster_changed():
     if config('prefer-ipv6') and rdata.get('hostname'):
         private_address = rdata['private-address']
         hostname = rdata['hostname']
-        rabbit.render_hosts(private_address, hostname)
+        if hostname:
+            rabbit.update_hosts_file({private_address: hostname})
 
     # sync passwords
-    peer_echo(['cookie', 'password'])
+    peer_echo(includes=['cookie', 'password'])
 
     # sync cookie
     cookie = peer_retrieve('cookie')

unit_tests/test_rabbit_utils.py

@@ -0,0 +1,71 @@
+import mock
+import os
+import unittest
+import tempfile
+import sys
+
+sys.modules['MySQLdb'] = mock.Mock()
+import rabbit_utils
+
+
+class UtilsTests(unittest.TestCase):
+    def setUp(self):
+        super(UtilsTests, self).setUp()
+
+    @mock.patch("rabbit_utils.log")
+    def test_update_empty_hosts_file(self, mock_log):
+        map = {'1.2.3.4': 'my-host'}
+        with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
+            rabbit_utils.HOSTS_FILE = tmpfile.name
+            rabbit_utils.HOSTS_FILE = tmpfile.name
+            rabbit_utils.update_hosts_file(map)
+
+        with open(tmpfile.name, 'r') as fd:
+            lines = fd.readlines()
+
+        os.remove(tmpfile.name)
+        self.assertEqual(len(lines), 1)
+        self.assertEqual(lines[0], "%s %s\n" % (map.items()[0]))
+
+    @mock.patch("rabbit_utils.log")
+    def test_update_hosts_file_w_dup(self, mock_log):
+        map = {'1.2.3.4': 'my-host'}
+        with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
+            rabbit_utils.HOSTS_FILE = tmpfile.name
+
+            with open(tmpfile.name, 'w') as fd:
+                fd.write("%s %s\n" % (map.items()[0]))
+
+            rabbit_utils.update_hosts_file(map)
+
+        with open(tmpfile.name, 'r') as fd:
+            lines = fd.readlines()
+
+        os.remove(tmpfile.name)
+        self.assertEqual(len(lines), 1)
+        self.assertEqual(lines[0], "%s %s\n" % (map.items()[0]))
+
+    @mock.patch("rabbit_utils.log")
+    def test_update_hosts_file_entry(self, mock_log):
+        altmap = {'1.1.1.1': 'alt-host'}
+        map = {'1.1.1.1': 'hostA',
+               '2.2.2.2': 'hostB',
+               '3.3.3.3': 'hostC',
+               '4.4.4.4': 'hostD'}
+        with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
+            rabbit_utils.HOSTS_FILE = tmpfile.name
+
+            with open(tmpfile.name, 'w') as fd:
+                fd.write("#somedata\n")
+                fd.write("%s %s\n" % (altmap.items()[0]))
+
+            rabbit_utils.update_hosts_file(map)
+
+        with open(rabbit_utils.HOSTS_FILE, 'r') as fd:
+            lines = fd.readlines()
+
+        os.remove(tmpfile.name)
+        self.assertEqual(len(lines), 5)
+        self.assertEqual(lines[0], "#somedata\n")
+        self.assertEqual(lines[1], "%s %s\n" % (map.items()[0]))
+        self.assertEqual(lines[4], "%s %s\n" % (map.items()[3]))