From e824e8174b3b479788a35cf537301cc3aed3ce2c Mon Sep 17 00:00:00 2001
From: Haowei Cai <haoweic@google.com>
Date: Thu, 26 Apr 2018 13:49:44 -0700
Subject: [PATCH] Watch properly decode resourceVersion from custom object
 response

---
 watch/watch.py      |  5 +++++
 watch/watch_test.py | 29 +++++++++++++++++++++++------
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/watch/watch.py b/watch/watch.py
index 6f10e7b5..91c4f924 100644
--- a/watch/watch.py
+++ b/watch/watch.py
@@ -84,6 +84,11 @@ def unmarshal_event(self, data, return_type):
             js['object'] = self._api_client.deserialize(obj, return_type)
             if hasattr(js['object'], 'metadata'):
                 self.resource_version = js['object'].metadata.resource_version
+            # For custom objects that we don't have model defined, json
+            # deserialization results in dictionary
+            elif isinstance(js['object'], dict) and 'metadata' in js['object']:
+                self.resource_version = js['object']['metadata'][
+                    'resourceVersion']
         return js
 
     def stream(self, func, *args, **kwargs):
diff --git a/watch/watch_test.py b/watch/watch_test.py
index 73bcc941..d1ec80a1 100644
--- a/watch/watch_test.py
+++ b/watch/watch_test.py
@@ -27,13 +27,13 @@ def test_watch_with_decode(self):
         fake_resp.release_conn = Mock()
         fake_resp.read_chunked = Mock(
             return_value=[
-                '{"type": "ADDED", "object": {"metadata": {"name": "test1"}'
-                ',"spec": {}, "status": {}}}\n',
-                '{"type": "ADDED", "object": {"metadata": {"name": "test2"}'
-                ',"spec": {}, "sta',
+                '{"type": "ADDED", "object": {"metadata": {"name": "test1",'
+                '"resourceVersion": "1"}, "spec": {}, "status": {}}}\n',
+                '{"type": "ADDED", "object": {"metadata": {"name": "test2",'
+                '"resourceVersion": "2"}, "spec": {}, "sta',
                 'tus": {}}}\n'
-                '{"type": "ADDED", "object": {"metadata": {"name": "test3"},'
-                '"spec": {}, "status": {}}}\n',
+                '{"type": "ADDED", "object": {"metadata": {"name": "test3",'
+                '"resourceVersion": "3"}, "spec": {}, "status": {}}}\n',
                 'should_not_happened\n'])
 
         fake_api = Mock()
@@ -46,6 +46,10 @@ def test_watch_with_decode(self):
             self.assertEqual("ADDED", e['type'])
             # make sure decoder worked and we got a model with the right name
             self.assertEqual("test%d" % count, e['object'].metadata.name)
+            # make sure decoder worked and updated Watch.resource_version
+            self.assertEqual(
+                "%d" % count, e['object'].metadata.resource_version)
+            self.assertEqual("%d" % count, w.resource_version)
             count += 1
             # make sure we can stop the watch and the last event with won't be
             # returned
@@ -133,6 +137,19 @@ def test_unmarshal_with_no_return_type(self):
         self.assertEqual(["test1"], event['object'])
         self.assertEqual(["test1"], event['raw_object'])
 
+    def test_unmarshal_with_custom_object(self):
+        w = Watch()
+        event = w.unmarshal_event('{"type": "ADDED", "object": {"apiVersion":'
+                                  '"test.com/v1beta1","kind":"foo","metadata":'
+                                  '{"name": "bar", "resourceVersion": "1"}}}',
+                                  'object')
+        self.assertEqual("ADDED", event['type'])
+        # make sure decoder deserialized json into dictionary and updated
+        # Watch.resource_version
+        self.assertTrue(isinstance(event['object'], dict))
+        self.assertEqual("1", event['object']['metadata']['resourceVersion'])
+        self.assertEqual("1", w.resource_version)
+
     def test_watch_with_exception(self):
         fake_resp = Mock()
         fake_resp.close = Mock()