[tor-commits] [stem/master] Add basic socket functions to Relay class

[atagar at torproject.org]

commit ad47e6f72cad9ceb61e7ff1dfdc4d8e7755aa0a8
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Feb 4 12:31:07 2018 -0800

    Add basic socket functions to Relay class
    
    Like the BaseController, callers will want basic socket passthroughs. Adding
    is_alive(), connection_time(), close(), and context management.
---
 stem/control.py                 |  7 +++----
 stem/relay.py                   | 40 ++++++++++++++++++++++++++++++++++++++++
 stem/socket.py                  |  1 +
 test/integ/client/connection.py | 17 +++++++++++++++++
 4 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/stem/control.py b/stem/control.py
index 123e3b20..b78707c1 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -153,8 +153,7 @@ If you're fine with allowing your script to raise exceptions then this can be mo
     |- get_socket - provides the socket used for control communication
     |- get_latest_heartbeat - timestamp for when we last heard from tor
     |- add_status_listener - notifies a callback of changes in our status
-    |- remove_status_listener - prevents further notification of status changes
-    +- __enter__ / __exit__ - manages socket connection
+    +- remove_status_listener - prevents further notification of status changes
 
 .. data:: State (enum)
 
@@ -674,7 +673,7 @@ class BaseController(object):
   def is_alive(self):
     """
     Checks if our socket is currently connected. This is a pass-through for our
-    socket's :func:`~stem.socket.ControlSocket.is_alive` method.
+    socket's :func:`~stem.socket.BaseSocket.is_alive` method.
 
     :returns: **bool** that's **True** if our socket is connected and **False** otherwise
     """
@@ -729,7 +728,7 @@ class BaseController(object):
   def close(self):
     """
     Closes our socket connection. This is a pass-through for our socket's
-    :func:`~stem.socket.ControlSocket.close` method.
+    :func:`~stem.socket.BaseSocket.close` method.
     """
 
     self._socket.close()
diff --git a/stem/relay.py b/stem/relay.py
index 20341721..14b7833f 100644
--- a/stem/relay.py
+++ b/stem/relay.py
@@ -13,6 +13,10 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
 
   Relay - Connection with a tor relay's ORPort.
     | +- connect - Establishes a connection with a relay.
+    |
+    |- is_alive - reports if our connection is open or closed
+    |- connection_time - time when we last connected or disconnected
+    +- close - shuts down our connection
 """
 
 import stem
@@ -95,3 +99,39 @@ class Relay(object):
     conn.send(stem.client.cell.NetinfoCell(stem.client.Address(address, addr_type), []).pack(link_protocol))
 
     return Relay(conn, link_protocol)
+
+  def is_alive(self):
+    """
+    Checks if our socket is currently connected. This is a pass-through for our
+    socket's :func:`~stem.socket.BaseSocket.is_alive` method.
+
+    :returns: **bool** that's **True** if our socket is connected and **False** otherwise
+    """
+
+    return self._orport.is_alive()
+
+  def connection_time(self):
+    """
+    Provides the unix timestamp for when our socket was either connected or
+    disconnected. That is to say, the time we connected if we're currently
+    connected and the time we disconnected if we're not connected.
+
+    :returns: **float** for when we last connected or disconnected, zero if
+      we've never connected
+    """
+
+    return self._orport.connection_time()
+
+  def close(self):
+    """
+    Closes our socket connection. This is a pass-through for our socket's
+    :func:`~stem.socket.BaseSocket.close` method.
+    """
+
+    return self._orport.close()
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exit_type, value, traceback):
+    self.close()
diff --git a/stem/socket.py b/stem/socket.py
index 7a2ec9de..bde68cc9 100644
--- a/stem/socket.py
+++ b/stem/socket.py
@@ -60,6 +60,7 @@ Tor...
     |
     |- is_alive - reports if the socket is known to be closed
     |- is_localhost - returns if the socket is for the local system or not
+    |- connection_time - timestamp when socket last connected or disconnected
     |- connect - connects a new socket
     |- close - shuts down the socket
     +- __enter__ / __exit__ - manages socket connection
diff --git a/test/integ/client/connection.py b/test/integ/client/connection.py
index bb6ae762..f5399a71 100644
--- a/test/integ/client/connection.py
+++ b/test/integ/client/connection.py
@@ -2,6 +2,7 @@
 Integration tests for establishing a connection with tor's ORPort.
 """
 
+import time
 import unittest
 
 import stem
@@ -40,6 +41,22 @@ class TestConnection(unittest.TestCase):
     for link_protocol in (1, 2, 6, 20):
       self.assertRaisesRegexp(stem.SocketError, 'Unable to establish a common link protocol with 127.0.0.1:1113', Relay.connect, '127.0.0.1', test.runner.ORPORT, [link_protocol])
 
+  def test_connection_time(self):
+    """
+    Checks duration we've been connected.
+    """
+
+    before = time.time()
+
+    with Relay.connect('127.0.0.1', test.runner.ORPORT) as conn:
+      connection_time = conn.connection_time()
+      self.assertTrue(time.time() >= connection_time >= before)
+      time.sleep(0.02)
+      self.assertTrue(conn.is_alive())
+
+    self.assertFalse(conn.is_alive())
+    self.assertTrue(conn.connection_time() >= connection_time + 0.02)
+
   def test_established(self):
     """
     Successfully establish ORPort connection.