From 435637ad3a908859ce08f9438e373a87fe5a3429 Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Sat, 1 Aug 2020 14:48:09 +0530
Subject: [PATCH] Add test for `process_post_breach_telemetry()`

---
 .../telemetry/processing/post_breach.py       | 15 ++--
 .../telemetry/processing/test_post_breach.py  | 85 +++++++++++++++++++
 2 files changed, 93 insertions(+), 7 deletions(-)
 create mode 100644 monkey/monkey_island/cc/services/telemetry/processing/test_post_breach.py

diff --git a/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py b/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
index fb54522b6..367ca87d9 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
+++ b/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}})
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/test_post_breach.py b/monkey/monkey_island/cc/services/telemetry/processing/test_post_breach.py
new file mode 100644
index 000000000..0999e285e
--- /dev/null
+++ b/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