From 045785c36fe1f5f6e3d37200e16f7abc4daf47b6 Mon Sep 17 00:00:00 2001
From: Marcelo Moreira de Mello <tchello.mello@gmail.com>
Date: Wed, 22 Sep 2021 23:03:45 -0400
Subject: [PATCH] Refactored get_conn_type() method to use Enum

---
 awx/main/tasks.py          |  2 +-
 awx/main/utils/receptor.py | 17 ++++++++++-------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/awx/main/tasks.py b/awx/main/tasks.py
index ac65899121..0f65b877da 100644
--- a/awx/main/tasks.py
+++ b/awx/main/tasks.py
@@ -3049,7 +3049,7 @@ class AWXReceptorJob:
         _kw = {}
         if self.work_type == 'ansible-runner':
             _kw['node'] = self.task.instance.execution_node
-            use_stream_tls = True if get_conn_type(_kw['node'], receptor_ctl) == 2 else False
+            use_stream_tls = get_conn_type(_kw['node'], receptor_ctl).name == "STREAMTLS"
             _kw['tlsclient'] = get_tls_client(use_stream_tls)
 
         result = receptor_ctl.submit_work(worktype=self.work_type, payload=sockout.makefile('rb'), params=self.receptor_params, **_kw)
diff --git a/awx/main/utils/receptor.py b/awx/main/utils/receptor.py
index fb45018614..fa1cb8f4cc 100644
--- a/awx/main/utils/receptor.py
+++ b/awx/main/utils/receptor.py
@@ -4,12 +4,20 @@ import time
 
 from receptorctl.socket_interface import ReceptorControl
 
+from enum import Enum, unique
 
 logger = logging.getLogger('awx.main.utils.receptor')
 
 __RECEPTOR_CONF = '/etc/receptor/receptor.conf'
 
 
+@unique
+class ReceptorConnectionType(Enum):
+    DATAGRAM = 0
+    STREAM = 1
+    STREAMTLS = 2
+
+
 def get_receptor_sockfile():
     with open(__RECEPTOR_CONF, 'r') as f:
         data = yaml.safe_load(f)
@@ -47,20 +55,15 @@ def get_receptor_ctl():
 
 
 def get_conn_type(node_name, receptor_ctl):
-    """
-    ConnType 0: Datagram
-    ConnType 1: Stream
-    ConnType 2: StreamTLS
-    """
     all_nodes = receptor_ctl.simple_command("status").get('Advertisements', None)
     for node in all_nodes:
         if node.get('NodeID') == node_name:
-            return node.get('ConnType')
+            return ReceptorConnectionType(node.get('ConnType'))
 
 
 def worker_info(node_name, work_type='ansible-runner'):
     receptor_ctl = get_receptor_ctl()
-    use_stream_tls = True if get_conn_type(node_name, receptor_ctl) == 2 else False
+    use_stream_tls = get_conn_type(node_name, receptor_ctl).name == "STREAMTLS"
     transmit_start = time.time()
     error_list = []
     data = {'errors': error_list, 'transmit_timing': 0.0}