diff options
Diffstat (limited to 'websocket')
-rw-r--r-- | websocket/_core.py | 13 | ||||
-rw-r--r-- | websocket/tests/test_websocket.py | 13 |
2 files changed, 26 insertions, 0 deletions
diff --git a/websocket/_core.py b/websocket/_core.py index 7d73d76..4d7ddec 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -154,6 +154,19 @@ class WebSocket(object): else: self.lock = NoLock() + def __iter__(self): + """ + Allow iteration over websocket, implying sequential `recv` executions. + """ + while True: + yield self.recv() + + def __next__(self): + return self.recv() + + def next(self): + return self.__next__() + def fileno(self): return self.sock.fileno() diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index 33f7a63..81249ec 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -255,6 +255,19 @@ class WebSocketTest(unittest.TestCase): data = sock.recv() self.assertEqual(data, "Hello") + @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") + def testIter(self): + count = 2 + for rsvp in ws.create_connection('ws://stream.meetup.com/2/rsvps'): + count -= 1 + if count == 0: + break + + @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") + def testNext(self): + sock = ws.create_connection('ws://stream.meetup.com/2/rsvps') + self.assertEqual(str, type(next(sock))) + def testInternalRecvStrict(self): sock = ws.WebSocket() s = sock.sock = SockMock() |