From f8d2c0bc81f476bfdf8f646805603973ca6f3032 Mon Sep 17 00:00:00 2001
From: Karl Hobley <karl@torchbox.com>
Date: Tue, 11 Feb 2014 14:28:20 +0000
Subject: [PATCH] update_index command is now callable from unit tests

---
 .../wagtailsearch/backends/elasticsearch.py   |  4 ++-
 .../management/commands/update_index.py       | 35 +++++++++++++------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/wagtail/wagtailsearch/backends/elasticsearch.py b/wagtail/wagtailsearch/backends/elasticsearch.py
index 56fa2b04dd..00c1bda0e9 100644
--- a/wagtail/wagtailsearch/backends/elasticsearch.py
+++ b/wagtail/wagtailsearch/backends/elasticsearch.py
@@ -178,9 +178,11 @@ class ElasticSearch(BaseSearch):
             type_set[obj_type].append(obj.indexed_build_document())
 
         # Loop through each type and bulk add them
+        results = []
         for type_name, type_objects in type_set.items():
-            print type_name, len(type_objects)
+            results.append((type_name, len(type_objects)))
             self.es.bulk_index(self.es_index, type_name, type_objects)
+        return results
 
     def delete(self, obj):
         # Object must be a decendant of Indexed and be a django model
diff --git a/wagtail/wagtailsearch/management/commands/update_index.py b/wagtail/wagtailsearch/management/commands/update_index.py
index 7e901bf813..e9c6fdd96e 100644
--- a/wagtail/wagtailsearch/management/commands/update_index.py
+++ b/wagtail/wagtailsearch/management/commands/update_index.py
@@ -1,13 +1,19 @@
-from django.core.management.base import NoArgsCommand
+from django.core.management.base import BaseCommand
 from django.db import models
 
 from wagtail.wagtailsearch import Indexed, get_search_backend
 
 
-class Command(NoArgsCommand):
-    def handle_noargs(self, **options):
+class Command(BaseCommand):
+    def handle(self, backend='default', **options):
+        # Check of we need to be quiet
+        quiet = False
+        if 'quiet' in options and options['quiet']:
+            quiet = True
+
         # Print info
-        print "Getting object list"
+        if not quiet:
+            print "Getting object list"
 
         # Get list of indexed models
         indexed_models = [model for model in models.get_models() if issubclass(model, Indexed)]
@@ -46,21 +52,30 @@ class Command(NoArgsCommand):
                     object_set[key] = obj
 
         # Search backend
-        s = get_search_backend()
+        s = get_search_backend(backend=backend)
 
         # Reset the index
-        print "Reseting index"
+        if not quiet:
+            print "Reseting index"
         s.reset_index()
 
         # Add types
-        print "Adding types"
+        if not quiet:
+            print "Adding types"
         for model in indexed_models:
             s.add_type(model)
 
         # Add objects to index
-        print "Adding objects"
-        s.add_bulk(object_set.values())
+        if not quiet:
+            print "Adding objects"
+        results = s.add_bulk(object_set.values())
+
+        # Print results
+        if not quiet and results:
+            for result in results:
+                print result[0], result[1]
 
         # Refresh index
-        print "Refreshing index"
+        if not quiet:
+            print "Refreshing index"
         s.refresh_index()