Convert mode argument to valid structure in create_service1388-format-service-mode

Signed-off-by: Joffrey F <joffrey@docker.com>

6 files changed, 106 insertions, 7 deletions

diff --git a/docker/api/service.py b/docker/api/service.py
index 0d8421e..d2621e6 100644
--- a/docker/api/service.py
+++ b/docker/api/service.py
@@ -1,5 +1,6 @@
 import warnings
 from .. import auth, errors, utils
+from ..types import ServiceMode
 
 
 class ServiceApiMixin(object):
@@ -18,8 +19,8 @@ class ServiceApiMixin(object):
             name (string): User-defined name for the service. Optional.
             labels (dict): A map of labels to associate with the service.
                 Optional.
-            mode (string): Scheduling mode for the service (``replicated`` or
-                ``global``). Defaults to ``replicated``.
+            mode (ServiceMode): Scheduling mode for the service (replicated
+                or global). Defaults to replicated.
             update_config (UpdateConfig): Specification for the update strategy
                 of the service. Default: ``None``
             networks (:py:class:`list`): List of network names or IDs to attach
@@ -49,6 +50,9 @@ class ServiceApiMixin(object):
             raise errors.DockerException(
                 'Missing mandatory Image key in ContainerSpec'
             )
+        if mode and not isinstance(mode, dict):
+            mode = ServiceMode(mode)
+
         registry, repo_name = auth.resolve_repository_name(image)
         auth_header = auth.get_config_header(self, registry)
         if auth_header:
@@ -191,8 +195,8 @@ class ServiceApiMixin(object):
             name (string): New name for the service. Optional.
             labels (dict): A map of labels to associate with the service.
                 Optional.
-            mode (string): Scheduling mode for the service (``replicated`` or
-                ``global``). Defaults to ``replicated``.
+            mode (ServiceMode): Scheduling mode for the service (replicated
+                or global). Defaults to replicated.
             update_config (UpdateConfig): Specification for the update strategy
                 of the service. Default: ``None``.
             networks (:py:class:`list`): List of network names or IDs to attach
@@ -222,6 +226,8 @@ class ServiceApiMixin(object):
         if labels is not None:
             data['Labels'] = labels
         if mode is not None:
+            if not isinstance(mode, dict):
+                mode = ServiceMode(mode)
             data['Mode'] = mode
         if task_template is not None:
             image = task_template.get('ContainerSpec', {}).get('Image', None)
diff --git a/docker/types/__init__.py b/docker/types/__init__.py
index 7230723..8e2fc17 100644
--- a/docker/types/__init__.py
+++ b/docker/types/__init__.py
@@ -4,6 +4,6 @@ from .healthcheck import Healthcheck
 from .networks import EndpointConfig, IPAMConfig, IPAMPool, NetworkingConfig
 from .services import (
     ContainerSpec, DriverConfig, EndpointSpec, Mount, Resources, RestartPolicy,
-    TaskTemplate, UpdateConfig
+    ServiceMode, TaskTemplate, UpdateConfig
 )
 from .swarm import SwarmSpec, SwarmExternalCA
diff --git a/docker/types/services.py b/docker/types/services.py
index 6e1ad32..ec0fcb1 100644
--- a/docker/types/services.py
+++ b/docker/types/services.py
@@ -348,3 +348,38 @@ def convert_service_ports(ports):
 
         result.append(port_spec)
     return result
+
+
+class ServiceMode(dict):
+    """
+        Indicate whether a service should be deployed as a replicated or global
+        service, and associated parameters
+
+        Args:
+            mode (string): Can be either ``replicated`` or ``global``
+            replicas (int): Number of replicas. For replicated services only.
+    """
+    def __init__(self, mode, replicas=None):
+        if mode not in ('replicated', 'global'):
+            raise errors.InvalidArgument(
+                'mode must be either "replicated" or "global"'
+            )
+        if mode != 'replicated' and replicas is not None:
+            raise errors.InvalidArgument(
+                'replicas can only be used for replicated mode'
+            )
+        self[mode] = {}
+        if replicas:
+            self[mode]['Replicas'] = replicas
+
+    @property
+    def mode(self):
+        if 'global' in self:
+            return 'global'
+        return 'replicated'
+
+    @property
+    def replicas(self):
+        if self.mode != 'replicated':
+            return None
+        return self['replicated'].get('Replicas')
diff --git a/docs/api.rst b/docs/api.rst
index 110b0a7..b5c1e92 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -110,5 +110,6 @@ Configuration types
 .. autoclass:: Mount
 .. autoclass:: Resources
 .. autoclass:: RestartPolicy
+.. autoclass:: ServiceMode
 .. autoclass:: TaskTemplate
 .. autoclass:: UpdateConfig
diff --git a/tests/integration/api_service_test.py b/tests/integration/api_service_test.py
index fc79400..77d7d28 100644
--- a/tests/integration/api_service_test.py
+++ b/tests/integration/api_service_test.py
@@ -251,3 +251,31 @@ class ServiceTest(BaseAPIIntegrationTest):
         con_spec = svc_info['Spec']['TaskTemplate']['ContainerSpec']
         assert 'Env' in con_spec
         assert con_spec['Env'] == ['DOCKER_PY_TEST=1']
+
+    def test_create_service_global_mode(self):
+        container_spec = docker.types.ContainerSpec(
+            'busybox', ['echo', 'hello']
+        )
+        task_tmpl = docker.types.TaskTemplate(container_spec)
+        name = self.get_service_name()
+        svc_id = self.client.create_service(
+            task_tmpl, name=name, mode='global'
+        )
+        svc_info = self.client.inspect_service(svc_id)
+        assert 'Mode' in svc_info['Spec']
+        assert 'Global' in svc_info['Spec']['Mode']
+
+    def test_create_service_replicated_mode(self):
+        container_spec = docker.types.ContainerSpec(
+            'busybox', ['echo', 'hello']
+        )
+        task_tmpl = docker.types.TaskTemplate(container_spec)
+        name = self.get_service_name()
+        svc_id = self.client.create_service(
+            task_tmpl, name=name,
+            mode=docker.types.ServiceMode('replicated', 5)
+        )
+        svc_info = self.client.inspect_service(svc_id)
+        assert 'Mode' in svc_info['Spec']
+        assert 'Replicated' in svc_info['Spec']['Mode']
+        assert svc_info['Spec']['Mode']['Replicated'] == {'Replicas': 5}
diff --git a/tests/unit/dockertypes_test.py b/tests/unit/dockertypes_test.py
index d11e4f0..5c470ff 100644
--- a/tests/unit/dockertypes_test.py
+++ b/tests/unit/dockertypes_test.py
@@ -7,7 +7,8 @@ import pytest
 from docker.constants import DEFAULT_DOCKER_API_VERSION
 from docker.errors import InvalidArgument, InvalidVersion
 from docker.types import (
-    EndpointConfig, HostConfig, IPAMConfig, IPAMPool, LogConfig, Mount, Ulimit,
+    EndpointConfig, HostConfig, IPAMConfig, IPAMPool, LogConfig, Mount,
+    ServiceMode, Ulimit,
 )
 
 try:
@@ -260,7 +261,35 @@ class IPAMConfigTest(unittest.TestCase):
         })
 
 
-class TestMounts(unittest.TestCase):
+class ServiceModeTest(unittest.TestCase):
+    def test_replicated_simple(self):
+        mode = ServiceMode('replicated')
+        assert mode == {'replicated': {}}
+        assert mode.mode == 'replicated'
+        assert mode.replicas is None
+
+    def test_global_simple(self):
+        mode = ServiceMode('global')
+        assert mode == {'global': {}}
+        assert mode.mode == 'global'
+        assert mode.replicas is None
+
+    def test_global_replicas_error(self):
+        with pytest.raises(InvalidArgument):
+            ServiceMode('global', 21)
+
+    def test_replicated_replicas(self):
+        mode = ServiceMode('replicated', 21)
+        assert mode == {'replicated': {'Replicas': 21}}
+        assert mode.mode == 'replicated'
+        assert mode.replicas == 21
+
+    def test_invalid_mode(self):
+        with pytest.raises(InvalidArgument):
+            ServiceMode('foobar')
+
+
+class MountTest(unittest.TestCase):
     def test_parse_mount_string_ro(self):
         mount = Mount.parse_mount_string("/foo/bar:/baz:ro")
         assert mount['Source'] == "/foo/bar"