fixes for #49 and #50

lomond/session.py

@@ -54,7 +54,7 @@ def __repr__(self):
         return "<ws-session '{}'>".format(self.websocket.url)
 
     @property
-    def _time(self):
+    def session_time(self):
         """Get the time since the socket started."""
         return time.time() - self._start_time
 
@@ -228,7 +228,7 @@ def _send_request(self):
 
     def _check_poll(self, poll):
         """Check if it is time for a poll."""
-        _time = self._time
+        _time = self.session_time
         if self._poll_start is None or _time - self._poll_start >= poll:
             self._poll_start = _time
             return True
@@ -239,7 +239,7 @@ def _check_auto_ping(self, ping_rate):
         """Check if a ping is required."""
         if not ping_rate:
             return
-        current_time = self._time
+        current_time = self.session_time
         if current_time > self._next_ping:
             # Calculate next ping time that is in the future.
             self._next_ping = (
@@ -253,7 +253,7 @@ def _check_auto_ping(self, ping_rate):
     def _check_ping_timeout(self, ping_timeout):
         """Check if the server is not responding to pings."""
         if ping_timeout:
-            time_since_last_pong = self._time - self._last_pong
+            time_since_last_pong = self.session_time - self._last_pong
             if time_since_last_pong > ping_timeout:
                 log.debug('ping_timeout time exceeded')
                 return True
@@ -262,10 +262,11 @@ def _check_ping_timeout(self, ping_timeout):
     def _check_close_timeout(self, close_timeout):
         """Check if the close timeout was tripped."""
         if not close_timeout:
-            return False
+            return
         sent_close_time = self.websocket.sent_close_time
-        if (sent_close_time is not None and 
-            self._time >= sent_close_time + close_timeout):
+        if sent_close_time is None:
+            return
+        if self.session_time >= sent_close_time + close_timeout:
             raise _ForceDisconnect(
                 "server didn't respond to close packet "
                 "within {}s".format(close_timeout)
@@ -301,7 +302,7 @@ def _regular(self, poll, ping_rate, ping_timeout, close_timeout):
             raise _ForceDisconnect(
                 'exceeded {:.0f}s ping timeout'.format(ping_timeout)
             )
-        self._check_close_timeout(close_timeout)    
+        self._check_close_timeout(close_timeout)
 
     def _send_pong(self, event):
         """Send a pong message in response to ping event."""
@@ -313,7 +314,7 @@ def _send_pong(self, event):
 
     def _on_pong(self, event):
         """Record last pong time."""
-        self._last_pong = self._time
+        self._last_pong = self.session_time
 
     def _on_ready(self):
         """Called when a ready event is received."""
@@ -347,7 +348,7 @@ def run(self,
         # Create socket and connect to remote server
         try:
             sock, proxy = self._connect()
-            self._sock = sock 
+            self._sock = sock
         except _SocketFail as error:
             yield events.ConnectFail('{}'.format(error))
             return
@@ -380,7 +381,7 @@ def _regular():
             return ()
 
         try:
-            while True:
+            while not websocket.is_closed:
                 readable = selector.wait_readable(poll)
                 for event in _regular():
                     yield event

lomond/websocket.py

@@ -202,7 +202,6 @@ def connect(self,
         )
         return run_generator
 
-
     def reset(self):
         """Reset the state."""
         self.state = self.State()
@@ -231,13 +230,10 @@ def close(self, code=None, reason=None):
         if self.is_closed:
             log.debug('%r already closed', self)
         else:
-            if code is None:
-                code = Status.NORMAL
-            if reason is None:
-                reason = b'goodbye'
-            self._send_close(code, reason)
-            self.state.closing = True
-            self.state.sent_close_time = time.time()
+            if not self.is_closing:
+                self._send_close(code, reason)
+                self.state.closing = True
+                self.state.sent_close_time = self.session.session_time
 
     def _on_close(self, message):
         """Close logic generator."""
@@ -469,7 +465,9 @@ def send_text(self, text):
 
     def _send_close(self, code, reason):
         """Send a close frame."""
-        frame_bytes = Frame.build_close_payload(code, reason)
+        _code = code if code is not None else Status.NORMAL
+        _reason = reason if reason is not None else b'goodbye'
+        frame_bytes = Frame.build_close_payload(_code, _reason)
         try:
             self.session.send(Opcode.CLOSE, frame_bytes)
         except (errors.WebSocketUnavailable, errors.TransportFail):

tests/test_live.py

@@ -1,5 +1,7 @@
 import lomond
 from lomond import events
+from lomond.session import WebsocketSession
+from lomond import selectors
 
 
 def test_echo():
@@ -74,7 +76,7 @@ def test_echo_poll():
             elif polls == 2:
                 # Covers some lesser used code paths
                 ws.state.closed = True
-                ws.session._sock.close()    
+                ws.session._sock.close()
 
 
 def test_not_ws():
@@ -86,6 +88,23 @@ def test_not_ws():
     assert events[1].name == 'connected'
     assert events[2].name == 'rejected'
     assert events[3].name == 'disconnected'
+    assert events[3].graceful
+
+
+class SelectSession(WebsocketSession):
+    _selector_cls = selectors.SelectSelector
+
+
+def test_not_ws_select():
+    """Test against a URL that doesn't serve websockets."""
+    ws = lomond.WebSocket('wss://www.google.com')
+    events = list(ws.connect(session_class=SelectSession))
+    assert len(events) == 4
+    assert events[0].name == 'connecting'
+    assert events[1].name == 'connected'
+    assert events[2].name == 'rejected'
+    assert events[3].name == 'disconnected'
+    assert events[3].graceful
 
 
 def test_no_url_wss():

tests/test_session.py

@@ -78,7 +78,7 @@ def pending(self):
 
 
 class FakeWebSocket(object):
-    
+
     sent_close_time = -100
 
     def send_pong(self, data):
@@ -426,7 +426,7 @@ def fake_recv(self):
 def test_on_pong(session):
     session._on_ready()
     session._on_pong(events.Pong(b'foo'))
-    assert session._time - session._last_pong < 0.01
+    assert session.session_time - session._last_pong < 0.01
 
 
 def test_context_manager():

tests/test_websocket.py

@@ -17,13 +17,20 @@ class FakeSession(object):
     def __init__(self, *args, **kwargs):
         self.socket_buffer = []
         self.run_called = False
+        self._t = 0.0
 
     def run(self, *args, **kwargs):
         self.run_called = True
 
     def send(self, opcode, bytes):
         self.socket_buffer.append((opcode, bytes))
 
+    @property
+    def session_time(self):
+        _t = self._t
+        self._t += 1.0
+        return _t
+
     def close(self):
         pass