diff options
Diffstat (limited to 'app/test/fetch/db_fetch/mock_cursor.py')
-rw-r--r-- | app/test/fetch/db_fetch/mock_cursor.py | 24 |
1 files changed, 22 insertions, 2 deletions
diff --git a/app/test/fetch/db_fetch/mock_cursor.py b/app/test/fetch/db_fetch/mock_cursor.py index 71efd3b..10c67e1 100644 --- a/app/test/fetch/db_fetch/mock_cursor.py +++ b/app/test/fetch/db_fetch/mock_cursor.py @@ -7,19 +7,39 @@ # which accompanies this distribution, and is available at # # http://www.apache.org/licenses/LICENSE-2.0 # ############################################################################### + + +def require_open(method): + def wrapped(self, *args, **kwargs): + if self.closed: + raise ValueError("Cursor is closed") + return method(self, *args, **kwargs) + return wrapped + + class MockCursor: def __init__(self, result): self.result = result self.current = 0 + self.closed = False + @require_open def __next__(self): if self.current < len(self.result): - next = self.result[self.current] + nxt = self.result[self.current] self.current += 1 - return next + return nxt else: raise StopIteration + @require_open def __iter__(self): return self + + @require_open + def fetchall(self): + return self.result + + def close(self): + self.closed = True |