From 24a6aed5baab342fc5a034109f700837d0d615a5 Mon Sep 17 00:00:00 2001
From: Michael Still <mikal@stillhq.com>
Date: Sat, 13 Oct 2012 10:41:37 +1100
Subject: [PATCH] Unit test replication_size.

Continuing to improve the unit test coverage of glance replicator.

Change-Id: I1ac5b6b2afe896b93a59526850dc617cce97eb64
---
 bin/glance-replicator                       |  6 ++--
 glance/tests/unit/test_glance_replicator.py | 31 +++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/bin/glance-replicator b/bin/glance-replicator
index bdf0d5f3ed..f8bbe60f2f 100755
--- a/bin/glance-replicator
+++ b/bin/glance-replicator
@@ -255,7 +255,7 @@ class ImageService(object):
         return headers, body
 
 
-def replication_size(options, args):
+def replication_size(options, args, imageservice=None):
     """%(prog)s size <server:port>
 
     Determine the size of a glance instance if dumped to disk.
@@ -269,7 +269,9 @@ def replication_size(options, args):
     total_size = 0
     count = 0
 
-    client = ImageService(httplib.HTTPConnection(server, port),
+    if not imageservice:
+        imageservice = ImageService
+    client = imageservice(httplib.HTTPConnection(server, port),
                           options.slavetoken)
     for image in client.get_images():
         logging.debug(_('Considering image: %(image)s') % {'image': image})
diff --git a/glance/tests/unit/test_glance_replicator.py b/glance/tests/unit/test_glance_replicator.py
index f1df6c597f..af216e228c 100644
--- a/glance/tests/unit/test_glance_replicator.py
+++ b/glance/tests/unit/test_glance_replicator.py
@@ -20,6 +20,7 @@ import json
 import os
 import StringIO
 import sys
+import UserDict
 
 from glance.tests import utils as test_utils
 
@@ -241,3 +242,33 @@ class ImageServiceTestCase(test_utils.BaseTestCase):
         c.conn.prime_request('PUT', 'v1/images/%s' % image_meta['id'],
                              '', image_meta_headers, 200, '', '')
         headers, body = c.add_image_meta(image_meta)
+
+
+class FakeImageService(object):
+    def __init__(self, http_conn, authtoken):
+        pass
+
+    def get_images(self):
+        return [{'status': 'active', 'size': 100, 'id': 123},
+                {'status': 'deleted', 'size': 200, 'id': 456},
+                {'status': 'active', 'size': 300, 'id': 789}]
+
+
+class ReplicationCommandsTestCase(test_utils.BaseTestCase):
+    def test_replication_size(self):
+        options = UserDict.UserDict()
+        options.slavetoken = 'slavetoken'
+        args = ['localhost:9292']
+
+        stdout = sys.stdout
+        sys.stdout = StringIO.StringIO()
+        try:
+            glance_replicator.replication_size(options, args,
+                                               imageservice=FakeImageService)
+            sys.stdout.seek(0)
+            output = sys.stdout.read()
+        finally:
+            sys.stdout = stdout
+
+        output = output.rstrip()
+        self.assertEqual(output, 'Total size is 400 bytes across 2 images')