Add test for `process_post_breach_telemetry()`

monkey/monkey_island/cc/services/telemetry/processing/post_breach.py

@@ -36,12 +36,6 @@ def process_post_breach_telemetry(telemetry_json):
             modified_data.append(temp)
         return modified_data
 
-    def update_data(data):
-        data = add_message_for_blank_outputs(data)
-        mongo.db.monkey.update(
-            {'guid': telemetry_json['monkey_guid']},
-            {'$push': {'pba_results': data}})
-
     def add_message_for_blank_outputs(data):
         if not data['result'][0]:
             data['result'][0] = EXECUTION_WITHOUT_OUTPUT
@@ -54,4 +48,11 @@ def process_post_breach_telemetry(telemetry_json):
     telemetry_json['data'] = convert_telem_data_to_list(telemetry_json['data'])
 
     for pba_data in telemetry_json['data']:
-        update_data(pba_data)
+        pba_data = add_message_for_blank_outputs(pba_data)
+        update_data(telemetry_json, pba_data)
+
+
+def update_data(telemetry_json, data):
+    mongo.db.monkey.update(
+        {'guid': telemetry_json['monkey_guid']},
+        {'$push': {'pba_results': data}})

monkey/monkey_island/cc/services/telemetry/processing/test_post_breach.py

@@ -0,0 +1,85 @@
+from unittest.mock import Mock
+
+import monkey_island.cc.services.telemetry.processing.post_breach as post_breach
+
+from .post_breach import EXECUTION_WITHOUT_OUTPUT
+
+original_telem_multiple_results =\
+    {
+        'data': {
+            'command': 'COMMAND',
+            'hostname': 'HOST',
+            'ip': '127.0.1.1',
+            'name': 'PBA NAME',
+            'result': [
+                ['SUCCESSFUL', True],
+                ['UNSUCCESFUL', False],
+                ['', True]
+            ]
+        },
+        'telem_category': 'post_breach'
+    }
+
+expected_telem_multiple_results =\
+    {
+        'data': [
+            {
+                'command': 'COMMAND',
+                'hostname': 'HOST',
+                'ip': '127.0.1.1',
+                'name': 'PBA NAME',
+                'result': ['SUCCESSFUL', True]
+            },
+            {
+                'command': 'COMMAND',
+                'hostname': 'HOST',
+                'ip': '127.0.1.1',
+                'name': 'PBA NAME',
+                'result': ['UNSUCCESFUL', False]
+            },
+            {
+                'command': 'COMMAND',
+                'hostname': 'HOST',
+                'ip': '127.0.1.1',
+                'name': 'PBA NAME',
+                'result': [EXECUTION_WITHOUT_OUTPUT, True]
+            }
+        ],
+        'telem_category': 'post_breach'
+    }
+
+original_telem_single_result =\
+    {
+        'data': {
+            'command': 'COMMAND',
+            'hostname': 'HOST',
+            'ip': '127.0.1.1',
+            'name': 'PBA NAME',
+            'result': ['', True]
+        },
+        'telem_category': 'post_breach'
+    }
+
+expected_telem_single_result =\
+    {
+        'data': [
+            {
+                'command': 'COMMAND',
+                'hostname': 'HOST',
+                'ip': '127.0.1.1',
+                'name': 'PBA NAME',
+                'result': [EXECUTION_WITHOUT_OUTPUT, True]
+            },
+        ],
+        'telem_category': 'post_breach'
+    }
+
+
+def test_process_post_breach_telemetry():
+    post_breach.update_data = Mock()  # actual behavior of update_data() is to access mongodb
+    # multiple results in PBA
+    post_breach.process_post_breach_telemetry(original_telem_multiple_results)
+    assert original_telem_multiple_results == expected_telem_multiple_results
+    # single result in PBA
+    post_breach.process_post_breach_telemetry(original_telem_single_result)
+    assert original_telem_single_result == expected_telem_single_result