Merge pull request #1846 from feliperuhland/issue-1734-scale-method

Added scale method to the Service model.

2 files changed, 67 insertions, 2 deletions

diff --git a/docker/models/services.py b/docker/models/services.py
index 009e455..337ed44 100644
--- a/docker/models/services.py
+++ b/docker/models/services.py
@@ -1,6 +1,6 @@
 import copy
-from docker.errors import create_unexpected_kwargs_error
-from docker.types import TaskTemplate, ContainerSpec
+from docker.errors import create_unexpected_kwargs_error, InvalidArgument
+from docker.types import TaskTemplate, ContainerSpec, ServiceMode
 from .resource import Model, Collection
 
 
@@ -105,6 +105,25 @@ class Service(Model):
         )
         return self.client.api.service_logs(self.id, is_tty=is_tty, **kwargs)
 
+    def scale(self, replicas):
+        """
+        Scale service container.
+
+        Args:
+            replicas (int): The number of containers that should be running.
+
+        Returns:
+            ``True``if successful.
+        """
+
+        if 'Global' in self.attrs['Spec']['Mode'].keys():
+            raise InvalidArgument('Cannot scale a global container')
+
+        service_mode = ServiceMode('replicated', replicas)
+        return self.client.api.update_service(self.id, self.version,
+                                              service_mode,
+                                              fetch_current_spec=True)
+
 
 class ServiceCollection(Collection):
     """Services on the Docker server."""
diff --git a/tests/integration/models_services_test.py b/tests/integration/models_services_test.py
index b4df745..ce83428 100644
--- a/tests/integration/models_services_test.py
+++ b/tests/integration/models_services_test.py
@@ -4,6 +4,8 @@ import docker
 
 from .. import helpers
 from .base import TEST_API_VERSION
+from docker.errors import InvalidArgument
+from docker.types.services import ServiceMode
 
 
 class ServiceTest(unittest.TestCase):
@@ -227,6 +229,50 @@ class ServiceTest(unittest.TestCase):
         spec = service.attrs['Spec']['TaskTemplate']['ContainerSpec']
         assert spec.get('Command') == ['sleep', '300']
 
+    def test_scale_method_service(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+        service = client.services.create(
+            # create arguments
+            name=helpers.random_name(),
+            # ContainerSpec arguments
+            image="alpine",
+            command="sleep 300",
+        )
+        tasks = []
+        while len(tasks) == 0:
+            tasks = service.tasks()
+        assert len(tasks) == 1
+        service.scale(2)
+        while len(tasks) == 1:
+            tasks = service.tasks()
+        assert len(tasks) >= 2
+        # check that the container spec is not overridden with None
+        service.reload()
+        spec = service.attrs['Spec']['TaskTemplate']['ContainerSpec']
+        assert spec.get('Command') == ['sleep', '300']
+
+    def test_scale_method_global_service(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+        mode = ServiceMode('global')
+        service = client.services.create(
+            name=helpers.random_name(),
+            image="alpine",
+            command="sleep 300",
+            mode=mode
+        )
+        tasks = []
+        while len(tasks) == 0:
+            tasks = service.tasks()
+        assert len(tasks) == 1
+        with self.assertRaises(InvalidArgument,
+                               msg='Cannot scale a global container'):
+            service.scale(2)
+
+        assert len(tasks) == 1
+        service.reload()
+        spec = service.attrs['Spec']['TaskTemplate']['ContainerSpec']
+        assert spec.get('Command') == ['sleep', '300']
+
     @helpers.requires_api_version('1.25')
     def test_restart_service(self):
         client = docker.from_env(version=TEST_API_VERSION)