Repository: opensistemas-hub/osbrain
Pull request #177 (2017-07-25 07:24:12)

Closing connections enhancements
Various improvements on how to close a connection (socket).

codecov[bot]: # [Codecov](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=h1) Report
> Merging [#177](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=desc) into [master](https://codecov.io/gh/opensistemas-hub/osbrain/commit/539758770e4de478658bfcf65c177dd3021502c9?src=pr&el=desc) will **increase** coverage by `<.01%`.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177/graphs/tree.svg?height=150&width=650&src=pr&token=LfvMoNOrOc)](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #177      +/-   ##
==========================================
+ Coverage   98.73%   98.74%   +<.01%     
==========================================
  Files          25       25              
  Lines        3083     3104      +21     
  Branches      232      235       +3     
==========================================
+ Hits         3044     3065      +21     
  Misses         26       26              
  Partials       13       13
```


| [Impacted Files](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [osbrain/agent.py](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=tree#diff-b3NicmFpbi9hZ2VudC5weQ==) | `96.44% <100%> (+0.05%)` | :arrow_up: |
| [osbrain/tests/test\_agent.py](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=tree#diff-b3NicmFpbi90ZXN0cy90ZXN0X2FnZW50LnB5) | `100% <100%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=footer). Last update [5397587...7c66344](https://codecov.io/gh/opensistemas-hub/osbrain/pull/177?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

Flood1993: I will squash commits if this gets accepted. Since there is not much code, I think it is easier to review all changes directly instead of per commit @Peque .

Patch:
diff --git a/docs/source/considerations.rst b/docs/source/considerations.rst
index 7597afb..2d77413 100644
--- a/docs/source/considerations.rst
+++ b/docs/source/considerations.rst
@@ -26,6 +26,32 @@ For more information simply refer to the
 `ØMQ guide <http://zguide.zeromq.org/page:all>`_.
 
 
+Closing connections
+===================
+
+For closing a specific connection from an agent we need to call the
+``close()`` method, which takes the alias of the socket from the
+connection we want to close as a parameter.
+
+.. code-block:: python
+
+   agent.bind('PUB', alias='connection')
+   ...
+   agent.close('connection')
+
+There is also a ``close_all()`` method, which takes no parameters and
+will close all user-defined connections of the agent.
+
+Remember that the `linger` value from the osBrain configuration will
+be used for the actual `socket.close()` calls in both methods. For more
+information, simply refer to the
+`ØMQ guide <http://zguide.zeromq.org/page:all>`_.
+
+.. note:: Closing a connection within an agent will have no effect on
+   any possible agents at the other end of the connection. Remember to
+   manually close them as well if the connection is not going to be reused.
+
+
 Adding new methods
 ==================
 
diff --git a/osbrain/agent.py b/osbrain/agent.py
index 8e57f1d..b06ba99 100644
--- a/osbrain/agent.py
+++ b/osbrain/agent.py
@@ -1503,7 +1503,7 @@ class Agent():
             self._kill_now = True
 
     def kill(self):
-        self.close_sockets()
+        self.close_all()
         self._pyroDaemon.shutdown()
 
     def get_unique_external_zmq_sockets(self):
@@ -1537,13 +1537,48 @@ class Agent():
 
         return set(external_sockets)
 
-    def close_sockets(self):
+    def has_socket(self, alias):
+        """
+        Return whether the agent has the passed socket internally stored.
+        """
+        return alias in self.socket
+
+    def close(self, alias):
+        """
+        Close a socket given its alias and clear its entry from the
+        `Agent.socket` dictionary.
+        """
+        sock = self.socket[alias]
+
+        # Each socket might be pointed by different keys
+        entries_to_delete = []
+        for k, v in self.socket.items():
+            if v == sock:
+                entries_to_delete.append(k)
+
+        for entry in entries_to_delete:
+            del self.socket[entry]
+
+        sock.close(linger=get_linger())
+
+    def close_all(self):
         """
         Close all non-internal zmq sockets.
         """
+        # Each socket might be pointed by different keys
+        sockets_to_delete = []
         for sock in self.get_unique_external_zmq_sockets():
+            sockets_to_delete.append(sock)
             sock.close(linger=get_linger())
 
+        entries_to_delete = []
+        for k, v in self.socket.items():
+            if v in sockets_to_delete:
+                entries_to_delete.append(k)
+
+        for entry in entries_to_delete:
+            del self.socket[entry]
+
     def ping(self):
         """
         A test method to check the readiness of the agent. Used for testing


Test patch:
diff --git a/osbrain/tests/test_agent.py b/osbrain/tests/test_agent.py
index 9811338..3e40f10 100644
--- a/osbrain/tests/test_agent.py
+++ b/osbrain/tests/test_agent.py
@@ -260,7 +260,7 @@ def test_linger(nsproxy, linger, sleep_time, should_receive):
     assert agent_dies('puller', nsproxy)
 
     pusher.send('push', 'foo')
-    pusher.close_sockets()
+    pusher.close_all()
     # After this timeout, depending on linger value, 'foo' will no longer be
     # on queue to be sent
     time.sleep(sleep_time)
@@ -477,6 +477,41 @@ def test_invalid_handlers(nsproxy):
         agent.bind('REP', handler=1.234)
 
 
+def test_has_socket(nsproxy):
+    """
+    Test the has_socket method.
+    """
+    a0 = run_agent('a0')
+    assert a0.has_socket('_loopback_safe')
+    assert not a0.has_socket('pub')
+    a0.bind('PUB', alias='pub')
+    assert a0.has_socket('pub')
+
+
+def test_close(nsproxy):
+    """
+    Test that closing a socket removes it from the socket entry.
+    """
+    a0 = run_agent('a0')
+    a0.bind('PUB', alias='pub')
+
+    a0.close('pub')
+    assert not a0.has_socket('pub')
+
+
+def test_close_all(nsproxy):
+    """
+    Test that after a call to `close_all`, only those non-internal are
+    actually closed.
+    """
+    a0 = run_agent('a0')
+    a0.bind('PUB', alias='pub')
+
+    a0.close_all()
+    assert not a0.has_socket('pub')
+    assert a0.has_socket('_loopback_safe')
+
+
 def test_log_levels(nsproxy):
     """
     Test different log levels: info, warning, error and debug. Debug messages