Don't set future result in __await__

4 files changed, 95 insertions, 46 deletions

diff --git a/slack/task.py b/slack/task.py
index 74c5fc1..9dfc817 100644
--- a/slack/task.py
+++ b/slack/task.py
@@ -48,7 +48,7 @@ class Future(Awaitable[T]):
         else:
             self.id = future_id
         self._state: Literal["PENDING", "CANCELLED", "FINISHED"] = "PENDING"
-        self._result: Optional[T] = None
+        self._result: T
         self._exception: Optional[BaseException] = None
         self._cancel_message = None
         self._callbacks: List[Callable[[Self], object]] = []
@@ -56,15 +56,12 @@ class Future(Awaitable[T]):
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}('{self.id}')"
 
-    def __await__(self) -> Generator[Future[T], T, T]:
-        if self.cancelled():
-            raise self._make_cancelled_error()
-        result = yield self
-        if isinstance(result, BaseException):
-            self.set_exception(result)
-            raise result
-        self.set_result(result)
-        return result
+    def __await__(self) -> Generator[Future[T], None, T]:
+        if not self.done():
+            yield self  # This tells Task to wait for completion.
+        if not self.done():
+            raise RuntimeError("await wasn't used with future")
+        return self.result()  # May raise too.
 
     def _make_cancelled_error(self):
         if self._cancel_message is None:
diff --git a/tests/test_hook_process_hashtable.py b/tests/test_hook_process_hashtable.py
index 9d0db36..d0d683b 100644
--- a/tests/test_hook_process_hashtable.py
+++ b/tests/test_hook_process_hashtable.py
@@ -17,13 +17,14 @@ def test_hook_process_hashtable(mock_method: MagicMock):
     coroutine = hook_process_hashtable(command, options, timeout)
     future = coroutine.send(None)
     assert isinstance(future, FutureProcess)
+    future.set_result((command, 0, "out", "err"))
 
     mock_method.assert_called_once_with(
         command, options, timeout, get_callback_name(weechat_task_cb), future.id
     )
 
     with pytest.raises(StopIteration) as excinfo:
-        coroutine.send((command, 0, "out", "err"))
+        coroutine.send(None)
     assert excinfo.value.value == (command, 0, "out", "err")
 
 
@@ -33,18 +34,24 @@ def test_hook_process_hashtable_chunked(mock_method: MagicMock):
     options = {"option": "1"}
     timeout = 123
     coroutine = hook_process_hashtable(command, options, timeout)
-    future = coroutine.send(None)
-    assert isinstance(future, FutureProcess)
+    future_1 = coroutine.send(None)
+    assert isinstance(future_1, FutureProcess)
+    future_1.set_result((command, -1, "o1", "e1"))
 
     mock_method.assert_called_once_with(
-        command, options, timeout, get_callback_name(weechat_task_cb), future.id
+        command, options, timeout, get_callback_name(weechat_task_cb), future_1.id
     )
 
-    assert isinstance(coroutine.send((command, -1, "o1", "e1")), FutureProcess)
-    assert isinstance(coroutine.send((command, -1, "o2", "e2")), FutureProcess)
+    future_2 = coroutine.send(None)
+    assert isinstance(future_2, FutureProcess)
+    future_2.set_result((command, -1, "o2", "e2"))
+
+    future_3 = coroutine.send(None)
+    assert isinstance(future_3, FutureProcess)
+    future_3.set_result((command, 0, "o3", "e3"))
 
     with pytest.raises(StopIteration) as excinfo:
-        coroutine.send((command, 0, "o3", "e3"))
+        coroutine.send(None)
     assert excinfo.value.value == (command, 0, "o1o2o3", "e1e2e3")
 
 
@@ -54,11 +61,14 @@ def test_hook_process_hashtable_wait_on_max_file_descriptors(
 ):
     mock_available_file_descriptors.return_value = 0
     coroutine = hook_process_hashtable("", {}, 0)
-    future = coroutine.send(None)
-    assert isinstance(future, FutureTimer)
+    future_1 = coroutine.send(None)
+    assert isinstance(future_1, FutureTimer)
+    future_1.set_result((0,))
 
     mock_available_file_descriptors.return_value = 9
-    assert isinstance(coroutine.send((0,)), FutureTimer)
+    future_2 = coroutine.send(None)
+    assert isinstance(future_2, FutureTimer)
+    future_2.set_result((0,))
 
     mock_available_file_descriptors.return_value = 10
-    assert isinstance(coroutine.send((0,)), FutureProcess)
+    assert isinstance(coroutine.send(None), FutureProcess)
diff --git a/tests/test_http_request.py b/tests/test_http_request.py
index 3c74e1f..a9e941d 100644
--- a/tests/test_http_request.py
+++ b/tests/test_http_request.py
@@ -15,6 +15,7 @@ def test_http_request_success(mock_method: MagicMock):
     options = {"option": "1"}
     timeout = 123
     coroutine = http_request(url, options, timeout)
+
     future = coroutine.send(None)
     assert isinstance(future, FutureProcess)
 
@@ -28,19 +29,23 @@ def test_http_request_success(mock_method: MagicMock):
 
     response = "response"
     body = f"HTTP/2 200\r\n\r\n{response}"
+    future.set_result(("", 0, body, ""))
 
     with pytest.raises(StopIteration) as excinfo:
-        coroutine.send(("", 0, body, ""))
+        coroutine.send(None)
     assert excinfo.value.value == response
 
 
 def test_http_request_error_process_return_code():
     url = "http://example.com"
     coroutine = http_request(url, {}, 0, max_retries=0)
-    assert isinstance(coroutine.send(None), FutureProcess)
+
+    future = coroutine.send(None)
+    assert isinstance(future, FutureProcess)
+    future.set_result(("", -2, "", ""))
 
     with pytest.raises(HttpError) as excinfo:
-        coroutine.send(("", -2, "", ""))
+        coroutine.send(None)
 
     assert excinfo.value.url == url
     assert excinfo.value.return_code == -2
@@ -51,10 +56,13 @@ def test_http_request_error_process_return_code():
 def test_http_request_error_process_stderr():
     url = "http://example.com"
     coroutine = http_request(url, {}, 0, max_retries=0)
-    assert isinstance(coroutine.send(None), FutureProcess)
+
+    future = coroutine.send(None)
+    assert isinstance(future, FutureProcess)
+    future.set_result(("", 0, "", "err"))
 
     with pytest.raises(HttpError) as excinfo:
-        coroutine.send(("", 0, "", "err"))
+        coroutine.send(None)
 
     assert excinfo.value.url == url
     assert excinfo.value.return_code == 0
@@ -65,13 +73,16 @@ def test_http_request_error_process_stderr():
 def test_http_request_error_process_http():
     url = "http://example.com"
     coroutine = http_request(url, {}, 0, max_retries=0)
-    assert isinstance(coroutine.send(None), FutureProcess)
+
+    future = coroutine.send(None)
+    assert isinstance(future, FutureProcess)
 
     response = "response"
     body = f"HTTP/2 400\r\n\r\n{response}"
+    future.set_result(("", 0, body, ""))
 
     with pytest.raises(HttpError) as excinfo:
-        coroutine.send(("", 0, body, ""))
+        coroutine.send(None)
 
     assert excinfo.value.url == url
     assert excinfo.value.return_code == 0
@@ -82,31 +93,53 @@ def test_http_request_error_process_http():
 def test_http_request_error_retry_success():
     url = "http://example.com"
     coroutine = http_request(url, {}, 0, max_retries=2)
-    assert isinstance(coroutine.send(None), FutureProcess)
 
-    assert isinstance(coroutine.send(("", -2, "", "")), FutureTimer)
-    assert isinstance(coroutine.send((0,)), FutureProcess)
+    future_1 = coroutine.send(None)
+    assert isinstance(future_1, FutureProcess)
+    future_1.set_result(("", -2, "", ""))
+
+    future_2 = coroutine.send(None)
+    assert isinstance(future_2, FutureTimer)
+    future_2.set_result((0,))
+
+    future_3 = coroutine.send(None)
+    assert isinstance(future_3, FutureProcess)
 
     response = "response"
     body = f"HTTP/2 200\r\n\r\n{response}"
+    future_3.set_result(("", 0, body, ""))
 
     with pytest.raises(StopIteration) as excinfo:
-        coroutine.send(("", 0, body, ""))
+        coroutine.send(None)
     assert excinfo.value.value == response
 
 
 def test_http_request_error_retry_error():
     url = "http://example.com"
     coroutine = http_request(url, {}, 0, max_retries=2)
-    assert isinstance(coroutine.send(None), FutureProcess)
 
-    assert isinstance(coroutine.send(("", -2, "", "")), FutureTimer)
-    assert isinstance(coroutine.send((0,)), FutureProcess)
-    assert isinstance(coroutine.send(("", -2, "", "")), FutureTimer)
-    assert isinstance(coroutine.send((0,)), FutureProcess)
+    future_1 = coroutine.send(None)
+    assert isinstance(future_1, FutureProcess)
+    future_1.set_result(("", -2, "", ""))
+
+    future_2 = coroutine.send(None)
+    assert isinstance(future_2, FutureTimer)
+    future_2.set_result((0,))
+
+    future_3 = coroutine.send(None)
+    assert isinstance(future_3, FutureProcess)
+    future_3.set_result(("", -2, "", ""))
+
+    future_4 = coroutine.send(None)
+    assert isinstance(future_4, FutureTimer)
+    future_4.set_result((0,))
+
+    future_5 = coroutine.send(None)
+    assert isinstance(future_5, FutureProcess)
+    future_5.set_result(("", -2, "", ""))
 
     with pytest.raises(HttpError) as excinfo:
-        coroutine.send(("", -2, "", ""))
+        coroutine.send(None)
 
     assert excinfo.value.url == url
     assert excinfo.value.return_code == -2
@@ -135,9 +168,10 @@ def test_http_request_multiple_headers():
         .strip()
         .replace("\n", "\r\n")
     )
+    future.set_result(("", 0, body, ""))
 
     with pytest.raises(StopIteration) as excinfo:
-        coroutine.send(("", 0, body, ""))
+        coroutine.send(future)
     assert excinfo.value.value == response
 
 
@@ -145,20 +179,27 @@ def test_http_request_multiple_headers():
 def test_http_request_ratelimit(mock_method: MagicMock):
     url = "http://example.com"
     coroutine = http_request(url, {}, 0)
-    assert isinstance(coroutine.send(None), FutureProcess)
+
+    future_1 = coroutine.send(None)
+    assert isinstance(future_1, FutureProcess)
 
     body = "HTTP/2 429\r\nRetry-After: 12\r\n\r\n"
+    future_1.set_result(("", 0, body, ""))
 
-    timer = coroutine.send(("", 0, body, ""))
-    assert isinstance(timer, FutureTimer)
+    future_2 = coroutine.send(None)
+    assert isinstance(future_2, FutureTimer)
+    future_2.set_result((0,))
 
     mock_method.assert_called_once_with(
-        12000, 0, 1, get_callback_name(weechat_task_cb), timer.id
+        12000, 0, 1, get_callback_name(weechat_task_cb), future_2.id
     )
 
-    assert isinstance(coroutine.send((0,)), FutureProcess)
+    future_3 = coroutine.send(None)
+    assert isinstance(future_3, FutureProcess)
 
     response = "response"
+    future_3.set_result(("", 0, f"HTTP/2 200\r\n\r\n{response}", ""))
+
     with pytest.raises(StopIteration) as excinfo:
-        coroutine.send(("", 0, f"HTTP/2 200\r\n\r\n{response}", ""))
+        coroutine.send(None)
     assert excinfo.value.value == response
diff --git a/tests/test_sleep.py b/tests/test_sleep.py
index 893c108..35a162c 100644
--- a/tests/test_sleep.py
+++ b/tests/test_sleep.py
@@ -13,11 +13,12 @@ def test_sleep(mock_method: MagicMock):
     coroutine = sleep(milliseconds)
     future = coroutine.send(None)
     assert isinstance(future, FutureTimer)
+    future.set_result((0,))
 
     mock_method.assert_called_once_with(
         milliseconds, 0, 1, get_callback_name(weechat_task_cb), future.id
     )
 
     with pytest.raises(StopIteration) as excinfo:
-        coroutine.send((0,))
+        coroutine.send(None)
     assert excinfo.value.value == (0,)  # TODO: Will probably change to None