Update unit tests

airbnb · Dec 14, 2017 · cd7dd18 · cd7dd18
1 parent 09b91d8
commit cd7dd18
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 66 deletions.
diff --git a/tests/manage_test.py b/tests/manage_test.py
@@ -6,15 +6,12 @@
 import os
 import subprocess
 import sys
-import time
-import uuid
 from unittest import mock, TestCase
 
 import boto3
 from pyfakefs import fake_filesystem_unittest
 
 import manage
-from tests.rules.eicar_rule_test import EICAR_STRING
 
 
 def _mock_input(prompt: str) -> str:
@@ -359,61 +356,9 @@ def test_deploy(self, mock_analyze: mock.MagicMock, mock_apply: mock.MagicMock,
         mock_apply.assert_called_once()
         mock_analyze.assert_called_once()
 
-    @mock.patch.object(time, 'sleep', mock.MagicMock())
-    @mock.patch.object(boto3, 'resource')
-    @mock.patch.object(manage, 'print')
-    @mock.patch.object(manage, 'pprint', mock.MagicMock())
-    @mock.patch.object(uuid, 'uuid4', return_value='test-uuid')
-    def test_live_test(self, mock_uuid: mock.MagicMock, mock_print: mock.MagicMock,
-                       mock_resource: mock.MagicMock):
-        """Verify execution order for boto3 and print mock calls."""
-        self.manager.live_test()
-
-        mock_uuid.assert_called_once()
-
-        mock_resource.assert_has_calls([
-            mock.call('s3'),
-            mock.call().Bucket('test.prefix.binaryalert-binaries.us-test-1'),
-            mock.call().Bucket().put_object(
-                Body=bytes('{}'.format(EICAR_STRING), 'utf-8'),
-                Key='eicar_test_test-uuid.txt',
-                Metadata={'filepath': 'eicar_test_test-uuid.txt'}
-            ),
-            mock.call('dynamodb'),
-            mock.call().Table('test_prefix_binaryalert_matches'),
-            mock.call().Table().query(
-                Select='ALL_ATTRIBUTES',
-                Limit=1,
-                ConsistentRead=True,
-                ScanIndexForward=False,
-                KeyConditionExpression=mock.ANY,
-                FilterExpression=mock.ANY
-            )
-        ])
-
-        mock_resource.assert_has_calls([
-            mock.call().Table().delete_item(Key=mock.ANY),
-            mock.call().Bucket().delete_objects(
-                Delete={'Objects': [{'Key': 'eicar_test_test-uuid.txt'}]}
-            )
-        ])
-
-        mock_print.assert_has_calls([
-            mock.call(
-                'Uploading EICAR test file '
-                'S3:test.prefix.binaryalert-binaries.us-test-1:eicar_test_test-uuid.txt...'
-            ),
-            mock.call(
-                'EICAR test file uploaded! '
-                'Connecting to table DynamoDB:test_prefix_binaryalert_matches...'
-            ),
-            mock.call(
-                '\t[1/10] Querying DynamoDB table for the expected YARA match entry...'
-            ),
-            mock.call('\nSUCCESS: Expected DynamoDB entry for the EICAR file was found!\n'),
-            mock.call('\nRemoving DynamoDB EICAR entry...'),
-            mock.call('Removing EICAR test file from S3...'),
-            mock.call(
-                '\nLive test succeeded! Verify the alert was sent to your SNS subscription(s).'
-            )
-        ])
+    @mock.patch.object(manage.live_test, 'run', return_value=False)
+    def test_live_test(self, mock_live_test: mock.MagicMock):
+        """Live test wrapper raises TestFailureError if appropriate."""
+        with self.assertRaises(manage.TestFailureError):
+            self.manager.live_test()
+        mock_live_test.assert_called_once()
diff --git a/tests/rules/eicar_rule_test.py b/tests/rules/eicar_rule_test.py
@@ -6,27 +6,32 @@
 
 THIS_DIRECTORY = os.path.dirname(os.path.realpath(__file__))  # Directory containing this file.
 EICAR_RULE_FILE = os.path.join(THIS_DIRECTORY, '..', '..', 'rules', 'public', 'eicar.yara')
-EICAR_STRING = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*"
+EICAR_TXT_FILE = os.path.join(THIS_DIRECTORY, '..', 'files', 'eicar.txt')
 
 
 class EicarRuleTest(unittest.TestCase):
     """Verify that the EICAR rules file matches only the expected string."""
     def setUp(self):
         """Compile the EICAR YARA rule."""
+        with open(EICAR_TXT_FILE, 'r') as f:
+            self.eicar_string = f.read()
         self.eicar_rule = yara.compile(EICAR_RULE_FILE)
 
     def test_match_eicar_string(self):
         """Should match the exact EICAR string."""
-        self.assertEqual(1, len(self.eicar_rule.match(data=EICAR_STRING)))
+        self.assertEqual(1, len(self.eicar_rule.match(data=self.eicar_string)))
 
     def test_match_eicar_with_trailing_spaces(self):
         """Trailing whitespace is allowed after the EICAR string."""
-        self.assertEqual(1, len(self.eicar_rule.match(data='{}    \n\t'.format(EICAR_STRING))))
+        matches = self.eicar_rule.match(data='{}    \n\t'.format(self.eicar_string))
+        self.assertEqual(1, len(matches))
 
     def test_no_match_if_eicar_is_not_beginning(self):
         """No match if EICAR string is not the beginning of the file."""
-        self.assertEqual(0, len(self.eicar_rule.match(data='other-text {}'.format(EICAR_STRING))))
+        matches = self.eicar_rule.match(data='other-text {}'.format(self.eicar_string))
+        self.assertEqual(0, len(matches))
 
     def test_no_match_if_eicar_is_not_end(self):
         """No match if non-whitespace comes after the EICAR string."""
-        self.assertEqual(0, len(self.eicar_rule.match(data='{} other-text'.format(EICAR_STRING))))
+        matches = self.eicar_rule.match(data='{} other-text'.format(self.eicar_string))
+        self.assertEqual(0, len(matches))