Update to sync query tests

Changes
=======
* Better handling of timeout errors
* Update sync tests to handle various forms (normal, lazy, cancellable)

Author: thejcfactor

couchbase_analytics/common/errors.py

@@ -124,8 +124,8 @@ class TimeoutError(AnalyticsError):
     Indicates that a request was unable to complete prior to reaching the deadline specified for the reqest.
     """
 
-    def __init__(self, base: Optional[Exception] = None, message: Optional[str] = None) -> None:
-        super().__init__(base, message)
+    def __init__(self, cause: Optional[Exception] = None, message: Optional[str] = None) -> None:
+        super().__init__(cause, message)
 
     def __repr__(self) -> str:
         return super().__repr__()

couchbase_analytics/protocol/streaming.py

@@ -32,7 +32,7 @@
 from httpx import Response as HttpCoreResponse
 
 # TODO: errors?
-from couchbase_analytics.common.errors import AnalyticsError, InternalSDKError
+from couchbase_analytics.common.errors import AnalyticsError, InternalSDKError, TimeoutError
 from couchbase_analytics.common.core import (JsonStreamConfig,
                                              ParsedResult,
                                              ParsedResultType)
@@ -69,7 +69,7 @@ def wrapped_fn(self: HttpStreamingResponse) -> None:
                     if self._request_context.request_error is not None:
                         raise self._request_context.request_error from None
                     if self._request_context.timed_out:
-                        raise TimeoutError() from None
+                        raise TimeoutError(message='Request timeout.') from None
                     if self._request_context.cancelled:
                         raise CancelledError('Request was cancelled.') from None
                     raise InternalSDKError(ex) from None
@@ -112,6 +112,17 @@ def _finish_processing_stream(self) -> None:
         while not self._json_stream.token_stream_exhausted:
             self._json_stream.continue_parsing()
 
+    def _handle_iteration_abort(self) -> None:
+        self.close()
+        if self._request_context.cancelled:
+            print('Request was cancelled, closing stream.')
+            raise StopIteration
+        elif self._request_context.timed_out:
+            print('Request timed out, closing stream.')
+            raise TimeoutError(message='Request timeout.')
+        else:
+            raise StopIteration
+
     def _maybe_continue_to_process_stream(self) -> None:
         if not self._request_context.has_stage_completed:
             return
@@ -190,17 +201,17 @@ def get_next_row(self) -> Any:
         if not (hasattr(self, '_core_response')
                 and self._core_response is not None
                 and self._request_context.okay_to_iterate):
-            self.close()
-            raise StopIteration
+            self._handle_iteration_abort()
 
         self._maybe_continue_to_process_stream()
+        check_state = False
         while True:
-            if self._request_context.cancelled:
-                self.close()
-                raise StopIteration
-            # TODO: handle timeout
+            if check_state and not self._request_context.okay_to_iterate:
+                self._handle_iteration_abort()
+
             raw_response = self._json_stream.get_result(self._stream_config.queue_timeout)
             if raw_response is None:
+                check_state = True
                 continue
             if raw_response.result_type == ParsedResultType.ROW:
                 if raw_response.value is None: