From c6610ad367ce9342b35db600e78e708781d2a6a4 Mon Sep 17 00:00:00 2001
From: Sergey Vasilyev <nolar@nolar.info>
Date: Thu, 11 Jul 2019 13:04:23 +0200
Subject: [PATCH] Remember the streaming response object in watch-queries (#31)

---
 pykube/query.py     |  6 ++++++
 tests/test_watch.py | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 tests/test_watch.py

diff --git a/pykube/query.py b/pykube/query.py
index 07e548e..ba4232d 100644
--- a/pykube/query.py
+++ b/pykube/query.py
@@ -157,6 +157,7 @@ class WatchQuery(BaseQuery):
     def __init__(self, *args, **kwargs):
         self.resource_version = kwargs.pop("resource_version", None)
         super(WatchQuery, self).__init__(*args, **kwargs)
+        self._response = None
 
     def object_stream(self):
         params = {"watch": "true"}
@@ -172,6 +173,7 @@ def object_stream(self):
             kwargs["version"] = self.api_obj_class.version
         r = self.api.get(**kwargs)
         self.api.raise_for_status(r)
+        self._response = r
         WatchEvent = namedtuple("WatchEvent", "type object")
         for line in r.iter_lines():
             we = json.loads(line.decode("utf-8"))
@@ -180,6 +182,10 @@ def object_stream(self):
     def __iter__(self):
         return iter(self.object_stream())
 
+    @property
+    def response(self):
+        return self._response
+
 
 def as_selector(value):
     if isinstance(value, str):
diff --git a/tests/test_watch.py b/tests/test_watch.py
new file mode 100644
index 0000000..6575908
--- /dev/null
+++ b/tests/test_watch.py
@@ -0,0 +1,40 @@
+import json
+from unittest.mock import MagicMock
+
+import pytest
+
+from pykube import Pod
+from pykube.query import Query
+
+
+@pytest.fixture
+def api():
+    return MagicMock()
+
+
+def test_watch_response_exists(api):
+    stream = Query(api, Pod).watch()
+    assert hasattr(stream, 'response')
+    assert stream.response is None  # not yet executed
+
+
+def test_watch_response_is_readonly(api):
+    stream = Query(api, Pod).watch()
+    with pytest.raises(AttributeError):
+        stream.response = object()
+
+
+def test_watch_response_is_set_on_iter(api):
+    line1 = json.dumps({'type': 'ADDED', 'object': {}}).encode('utf-8')
+    expected_response = MagicMock()
+    expected_response.iter_lines.return_value = [line1]
+    api.get.return_value = expected_response
+
+    stream = Query(api, Pod).watch()
+    next(iter(stream))
+
+    assert stream.response is expected_response
+
+    assert api.get.call_count == 1
+    assert api.get.call_args_list[0][1]['stream'] is True
+    assert 'watch=true' in api.get.call_args_list[0][1]['url']