From 2b6c935f89b5ff2d0717e04890dbe295f8cd0ea0 Mon Sep 17 00:00:00 2001
From: Bryan Worrell <bworrell@users.noreply.github.com>
Date: Fri, 19 Feb 2021 13:17:39 -0500
Subject: [PATCH] Fix issue with deletion of sessions during refresh() (#2055)

* Fix issue with deletion of sessions during refresh().

* Remove use of IsolatedAsyncioTestCase.

It's not available in all versions of Python we support.

Co-authored-by: Bryan Worrell <bworrell@mitre.org>
---
 app/contacts/contact_tcp.py        |  9 +++++++-
 tests/contacts/test_contact_tcp.py | 37 ++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 tests/contacts/test_contact_tcp.py

diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py
index 34dc75456..c5010c861 100644
--- a/app/contacts/contact_tcp.py
+++ b/app/contacts/contact_tcp.py
@@ -47,11 +47,18 @@ def __init__(self, services, log):
         self.sessions = []
 
     async def refresh(self):
-        for index, session in enumerate(self.sessions):
+        index = 0
+
+        while index < len(self.sessions):
+            session = self.sessions[index]
+
             try:
                 session.connection.send(str.encode(' '))
             except socket.error:
+                self.log.debug('Error occurred when refreshing session %s. Removing from session pool.', session.id)
                 del self.sessions[index]
+            else:
+                index += 1
 
     async def accept(self, reader, writer):
         try:
diff --git a/tests/contacts/test_contact_tcp.py b/tests/contacts/test_contact_tcp.py
new file mode 100644
index 000000000..337d691ab
--- /dev/null
+++ b/tests/contacts/test_contact_tcp.py
@@ -0,0 +1,37 @@
+import logging
+import socket
+from unittest import mock
+
+from app.contacts.contact_tcp import TcpSessionHandler
+
+logger = logging.getLogger(__name__)
+
+
+class TestTcpSessionHandler:
+
+    def test_refresh_with_socket_errors(self, loop):
+        handler = TcpSessionHandler(services=None, log=logger)
+
+        session_with_socket_error = mock.Mock()
+        session_with_socket_error.connection.send.side_effect = socket.error()
+
+        handler.sessions = [
+            session_with_socket_error,
+            session_with_socket_error,
+            mock.Mock()
+        ]
+
+        loop.run_until_complete(handler.refresh())
+        assert len(handler.sessions) == 1
+        assert all(x is not session_with_socket_error for x in handler.sessions)
+
+    def test_refresh_without_socket_errors(self, loop):
+        handler = TcpSessionHandler(services=None, log=logger)
+        handler.sessions = [
+            mock.Mock(),
+            mock.Mock(),
+            mock.Mock()
+        ]
+
+        loop.run_until_complete(handler.refresh())
+        assert len(handler.sessions) == 3