From 2109ec06f46f90dcc0d4c56bd8b3506f8623bb2a Mon Sep 17 00:00:00 2001
From: Will Kahn-Greene <willkg@mozilla.com>
Date: Thu, 9 Oct 2014 12:45:40 -0400
Subject: [PATCH] Fix bulk monkeypatch to work with non-insert op_types

This also adds a test for the monkeypatch to make sure it works with the
various bulk op_types.

Fixes #266
---
 elasticutils/monkeypatch.py            |  7 ++-
 elasticutils/tests/test_monkeypatch.py | 85 ++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 elasticutils/tests/test_monkeypatch.py

diff --git a/elasticutils/monkeypatch.py b/elasticutils/monkeypatch.py
index f41b0ae..60beeb3 100644
--- a/elasticutils/monkeypatch.py
+++ b/elasticutils/monkeypatch.py
@@ -24,8 +24,11 @@ def monkeypatch_es():
         @wraps(fun)
         def _fixed_bulk(self, *args, **kwargs):
             def fix_item(item):
-                if 'ok' in item['index']:
-                    item['index']['status'] = 201
+                # Go through all the possible sections of item looking
+                # for 'ok' and adding an additional 'status'.
+                for key, val in item.items():
+                    if 'ok' in val:
+                        val['status'] = 201
                 return item
 
             ret = fun(self, *args, **kwargs)
diff --git a/elasticutils/tests/test_monkeypatch.py b/elasticutils/tests/test_monkeypatch.py
new file mode 100644
index 0000000..ffecf62
--- /dev/null
+++ b/elasticutils/tests/test_monkeypatch.py
@@ -0,0 +1,85 @@
+from nose.tools import eq_
+
+from elasticutils.tests import ESTestCase
+
+
+class MonkeyPatchTest(ESTestCase):
+    mapping_type_name = 'eutestcolor'
+        
+    def test_bulk_insert_update_delete(self):
+        es = self.get_es()
+
+        # Bulk index two things and then verify they made it into the
+        # index.
+        data = [
+            {'index': {'_index': self.index_name,
+                       '_type': self.mapping_type_name,
+                       '_id': 1}},
+            {'color': 'blue'},
+
+            {'index': {'_index': self.index_name,
+                       '_type': self.mapping_type_name,
+                       '_id': 2}},
+            {'color': 'red'},
+        ]
+        es.bulk(data, refresh=True)
+
+        eq_(len(self.get_s()), 2)
+        eq_(self.get_s().filter(color='blue')[0]._id, '1')
+        eq_(self.get_s().filter(color='red')[0]._id, '2')
+
+        # Then bulk update them.
+        data = [
+            {'update': {'_index': self.index_name,
+                        '_type': self.mapping_type_name,
+                        '_id': 1}},
+            {'doc': {'color': 'green'}},
+
+            {'update': {'_index': self.index_name,
+                        '_type': self.mapping_type_name,
+                        '_id': 2}},
+            {'doc': {'color': 'pink'}}
+        ]
+        es.bulk(data, refresh=True)
+            
+        eq_(len(self.get_s()), 2)
+        eq_(len(self.get_s().filter(color='blue')), 0)
+        eq_(len(self.get_s().filter(color='red')), 0)
+
+        eq_(self.get_s().filter(color='green')[0]._id, '1')
+        eq_(self.get_s().filter(color='pink')[0]._id, '2')
+
+        # Then delete them and make sure they're gone.
+        data = [
+            {'delete': {'_index': self.index_name,
+                        '_type': self.mapping_type_name,
+                        '_id': 1}},
+            {'delete': {'_index': self.index_name,
+                        '_type': self.mapping_type_name,
+                        '_id': 2}}
+        ]
+        es.bulk(data, refresh=True)
+
+        eq_(len(self.get_s()), 0)
+
+    def test_bulk_create(self):
+        es = self.get_es()
+
+        # Bulk create two things and then verify they made it into the
+        # index.
+        data = [
+            {'create': {'_index': self.index_name,
+                       '_type': self.mapping_type_name,
+                       '_id': 1}},
+            {'color': 'blue'},
+
+            {'create': {'_index': self.index_name,
+                       '_type': self.mapping_type_name,
+                       '_id': 2}},
+            {'color': 'red'},
+        ]
+        es.bulk(data, refresh=True)
+
+        eq_(len(self.get_s()), 2)
+        eq_(self.get_s().filter(color='blue')[0]._id, '1')
+        eq_(self.get_s().filter(color='red')[0]._id, '2')