| |
@@ -54,15 +54,15 @@
|
| |
self.verify = verify
|
| |
self.url = url
|
| |
|
| |
- def retrieve(self, subject_type, subject_identifier, testcase=None):
|
| |
+ def retrieve(self, subject_type, subject_identifier, testcase=None, limit=None):
|
| |
"""
|
| |
Return generator over results.
|
| |
"""
|
| |
- for result in self._retrieve_helper(subject_type, subject_identifier, testcase):
|
| |
+ for result in self._retrieve_helper(subject_type, subject_identifier, testcase, limit):
|
| |
if result['id'] not in self.ignore_results:
|
| |
yield result
|
| |
|
| |
- def _retrieve_helper(self, subject_type, subject_identifier, testcase):
|
| |
+ def _retrieve_helper(self, subject_type, subject_identifier, testcase, limit=None):
|
| |
cache_key = results_cache_key(
|
| |
subject_type, subject_identifier, testcase)
|
| |
|
| |
@@ -75,14 +75,18 @@
|
| |
|
| |
while cached_results.can_fetch_more:
|
| |
cached_results.last_page += 1
|
| |
- results = self._retrieve_page(
|
| |
- cached_results.last_page, subject_type, subject_identifier,
|
| |
- testcase)
|
| |
- cached_results.results.extend(results)
|
| |
- cached_results.can_fetch_more = bool(results)
|
| |
- self.cache.set(cache_key, cached_results)
|
| |
- for result in results:
|
| |
- yield result
|
| |
+ if limit is None or (limit is not None and cached_results.last_page < limit):
|
| |
+ results = self._retrieve_page(
|
| |
+ cached_results.last_page, subject_type, subject_identifier,
|
| |
+ testcase)
|
| |
+ cached_results.results.extend(results)
|
| |
+ cached_results.can_fetch_more = bool(results)
|
| |
+ self.cache.set(cache_key, cached_results)
|
| |
+ for result in results:
|
| |
+ yield result
|
| |
+ else:
|
| |
+ log.warning("The list of results was truncated. Too many items.")
|
| |
+ break
|
| |
|
| |
def _make_request(self, params):
|
| |
response = requests_session.get(
|
| |