Include tag in images.get after pulling if provided separately1585-pull-tag

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

4 files changed, 16 insertions, 9 deletions

diff --git a/docker/models/images.py b/docker/models/images.py
index 48f3590..a12ede3 100644
--- a/docker/models/images.py
+++ b/docker/models/images.py
@@ -238,7 +238,7 @@ class ImageCollection(Collection):
         """
         return self.client.api.load_image(data)
 
-    def pull(self, name, **kwargs):
+    def pull(self, name, tag=None, **kwargs):
         """
         Pull an image of the given name and return it. Similar to the
         ``docker pull`` command.
@@ -267,8 +267,8 @@ class ImageCollection(Collection):
 
             >>> image = client.images.pull('busybox')
         """
-        self.client.api.pull(name, **kwargs)
-        return self.get(name)
+        self.client.api.pull(name, tag=tag, **kwargs)
+        return self.get('{0}:{1}'.format(name, tag) if tag else name)
 
     def push(self, repository, tag=None, **kwargs):
         return self.client.api.push(repository, tag=tag, **kwargs)
diff --git a/tests/integration/models_images_test.py b/tests/integration/models_images_test.py
index 49e06f6..6d61e49 100644
--- a/tests/integration/models_images_test.py
+++ b/tests/integration/models_images_test.py
@@ -30,10 +30,12 @@ class ImageCollectionTest(BaseIntegrationTest):
 
     def test_build_with_multiple_success(self):
         client = docker.from_env(version=TEST_API_VERSION)
-        image = client.images.build(tag='some-tag', fileobj=io.BytesIO(
-            "FROM alpine\n"
-            "CMD echo hello world".encode('ascii')
-            ))
+        image = client.images.build(
+            tag='some-tag', fileobj=io.BytesIO(
+                "FROM alpine\n"
+                "CMD echo hello world".encode('ascii')
+            )
+        )
         self.tmp_imgs.append(image.id)
         assert client.containers.run(image) == b"hello world\n"
 
@@ -53,6 +55,11 @@ class ImageCollectionTest(BaseIntegrationTest):
         image = client.images.pull('alpine:latest')
         assert 'alpine:latest' in image.attrs['RepoTags']
 
+    def test_pull_with_tag(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+        image = client.images.pull('alpine', tag='3.3')
+        assert 'alpine:3.3' in image.attrs['RepoTags']
+
 
 class ImageTest(BaseIntegrationTest):
 
diff --git a/tests/unit/models_containers_test.py b/tests/unit/models_containers_test.py
index e74bb7c..0fb69f3 100644
--- a/tests/unit/models_containers_test.py
+++ b/tests/unit/models_containers_test.py
@@ -227,7 +227,7 @@ class ContainerCollectionTest(unittest.TestCase):
         container = client.containers.run('alpine', 'sleep 300', detach=True)
 
         assert container.id == FAKE_CONTAINER_ID
-        client.api.pull.assert_called_with('alpine')
+        client.api.pull.assert_called_with('alpine', tag=None)
 
     def test_run_with_error(self):
         client = make_fake_client()
diff --git a/tests/unit/models_images_test.py b/tests/unit/models_images_test.py
index 784717b..9ecb7e4 100644
--- a/tests/unit/models_images_test.py
+++ b/tests/unit/models_images_test.py
@@ -42,7 +42,7 @@ class ImageCollectionTest(unittest.TestCase):
     def test_pull(self):
         client = make_fake_client()
         image = client.images.pull('test_image')
-        client.api.pull.assert_called_with('test_image')
+        client.api.pull.assert_called_with('test_image', tag=None)
         client.api.inspect_image.assert_called_with('test_image')
         assert isinstance(image, Image)
         assert image.id == FAKE_IMAGE_ID