Fix when pyyaml has interpreted line as a dictionary

Added unit tests in build_extra_hosts + fix

Signed-off-by: CJ <[email protected]>

compose/service.py

@@ -629,7 +629,17 @@ def build_extra_hosts(extra_hosts_config):
     if extra_hosts_config is None:
         return None
 
-    if isinstance(extra_hosts_config, list):
-        return dict(r.split(':') for r in extra_hosts_config)
-    else:
-        return dict([extra_hosts_config.split(':')])
+    if isinstance(extra_hosts_config, basestring):
+        extra_hosts_config = [extra_hosts_config]
+
+    extra_hosts_dict = {}
+    for extra_hosts_line in extra_hosts_config:
+        if isinstance(extra_hosts_line, dict):
+            # already interpreted as a dict (depends on pyyaml version)
+            extra_hosts_dict.update(extra_hosts_line)
+        else:
+            # not already interpreted as a dict
+            host, ip = extra_hosts_line.split(':')
+            extra_hosts_dict.update({host.strip(): ip.strip()})
+
+    return extra_hosts_dict

tests/integration/service_test.py

@@ -6,6 +6,7 @@ import mock
 
 from compose import Service
 from compose.service import CannotBeScaledError
+from compose.service import build_extra_hosts
 from compose.container import Container
 from docker.errors import APIError
 from .testcases import DockerClientTestCase
@@ -107,6 +108,28 @@ class ServiceTest(DockerClientTestCase):
         service.start_container(container)
         self.assertEqual(container.inspect()['Config']['CpuShares'], 73)
 
+    def test_build_extra_hosts(self):
+        # string
+        self.assertEqual(build_extra_hosts("www.example.com: 192.168.0.17"),
+                         {'www.example.com': '192.168.0.17'})
+
+        # list of strings
+        self.assertEqual(build_extra_hosts(
+                          ["www.example.com: 192.168.0.17"]),
+                         {'www.example.com': '192.168.0.17'})
+        self.assertEqual(build_extra_hosts(
+                          ["www.example.com: 192.168.0.17",
+                           "api.example.com: 192.168.0.18"]),
+                         {'www.example.com': '192.168.0.17',
+                          'api.example.com': '192.168.0.18'})
+        # list of dictionaries
+        self.assertEqual(build_extra_hosts(
+                          [{'www.example.com': '192.168.0.17'},
+                           {'api.example.com': '192.168.0.18'}
+                          ]),
+                          {'www.example.com': '192.168.0.17',
+                           'api.example.com': '192.168.0.18'})
+
     def test_create_container_with_extra_hosts_list(self):
         extra_hosts = ['docker:162.242.195.82', 'fig:50.31.209.229']
         service = self.create_service('db', extra_hosts=extra_hosts)