import copy from datetime import datetime from unittest.mock import patch, MagicMock import pytest from bson.objectid import ObjectId from CommonServerPython import DemistoException from MongoDB import ( Client, bulk_update_command, convert_id_to_object_id, convert_object_id_to_str, convert_str_to_datetime, format_sort, parse_and_validate_bulk_update_arguments, pipeline_query_command, search_query, update_entry_command, ) id_to_obj_inputs = [ ( [{"_id": "5e4412f230c5b8f63a7356ba"}], [{"_id": ObjectId("5e4412f230c5b8f63a7356ba")}], ), ( {"_id": "5e4412f230c5b8f63a7356ba"}, {"_id": ObjectId("5e4412f230c5b8f63a7356ba")}, ), ( {"_id": {"$gte": "5e4412f230c5b8f63a7356ba"}}, {"_id": {"$gte": ObjectId("5e4412f230c5b8f63a7356ba")}}, ), ({}, {}), ({"id": 1}, {"id": 1}), ] @pytest.mark.parametrize("func_input, expected", id_to_obj_inputs) def test_convert_id_to_object_id(func_input, expected): assert expected == convert_id_to_object_id(func_input) object_to_id = [ ([ObjectId("5e4412f230c5b8f63a7356ba")], ["5e4412f230c5b8f63a7356ba"]), ( [ObjectId("5e4412f230c5b8f63a7356ba"), ObjectId("5e4412f230c5b8f63a7356ba")], ["5e4412f230c5b8f63a7356ba", "5e4412f230c5b8f63a7356ba"], ), ] @pytest.mark.parametrize("func_input, expected", object_to_id) def test_convert_object_id_to_str(func_input, expected): assert expected == convert_object_id_to_str(func_input) def test_normalize_id(): res = Client.normalize_id({"_id": ObjectId("5e4412f230c5b8f63a7356ba")}) assert res["_id"] == "5e4412f230c5b8f63a7356ba" class TestConvertStrToDatetime: dict_inputs = [ {"testing": 123, "time": "ISODate('2020-06-12T08:23:07.000Z')"}, pytest.param({"testing": 123, "time": "ISODate('2018-06-12T08:23:07.000')"}, marks=pytest.mark.xfail), ] @pytest.mark.parametrize("func_input", dict_inputs) def test_convert_str_to_datetime(self, func_input): res = convert_str_to_datetime(func_input) assert isinstance(res["time"], datetime) def test_convert_str_to_datetime_no_datetime_obj(self): inputs = {1: 2} res = convert_str_to_datetime(inputs) assert isinstance(res[1], int) def test_nested_dict(self): """ Given: A nested dict with a timestamp When: Running a query or insert Then: Validating all keys in the dict are there and the timestamp is valid """ func_input = {"k": {"$gte": "ISODate('2020-06-12T08:23:07.000Z')"}} res = convert_str_to_datetime(func_input) assert isinstance(res["k"]["$gte"], datetime) class TestDatetimeToStr: datetime_obj = datetime.strptime("2020-05-19T09:05:28.000Z", "%Y-%m-%dT%H:%M:%S.000Z") datetime_str = "2020-05-19T09:05:28.000Z" def test_datetime_to_str_dict(self): """ Given: dict containing datetime object When: converting datetimes to strs Then: validate the value is a string. """ raw = Client.datetime_to_str({"time": self.datetime_obj}) assert self.datetime_str == raw["time"] def test_datetime_to_str_list(self): """ Given: list containing datetime object When: converting datetimes to strs Then: validate the value is a string. """ raw = Client.datetime_to_str([self.datetime_obj]) assert [self.datetime_str] == raw def test_datetime_to_str_str(self): """ Given: datetime object When: converting datetimes to strs Then: validate the value is a string. """ raw = Client.datetime_to_str(self.datetime_obj) assert self.datetime_str == raw def test_datetime_to_str_dict_no_datetime(self): """ Given: dict containing 5 (int) When: converting datetimes to strs Then: validate the value returned is 5 """ raw = Client.datetime_to_str({"time": 5}) assert raw["time"] == 5 def test_datetime_to_str_list_no_datetime(self): """ Given: list containing an int (5) object When: converting datetimes to strs Then: validate the value returned is 5. """ raw = Client.datetime_to_str([5]) assert raw == [5] def test_datetime_to_str_str_no_datetime(self): """ Given: 'str' When: converting datetimes to strs Then: validate the value returned is 'str'. """ raw = Client.datetime_to_str("str") assert raw == "str" class MockedQuery: class Limit: @staticmethod def limit(number): return [{"time": TestDatetimeToStr.datetime_obj, "_id": ObjectId("5e4412f230c5b8f63a7356ba")}] @classmethod def find(cls, query): return cls.Limit def test_query(mocker): """ Given: Object with datetime and object id in it When: Quering object Then: validate all objects returned are strs. """ client = Client(["aaaaa"], "a", "b", "d") mocker.patch.object(Client, "get_collection", return_value=MockedQuery) readable_outputs, outputs, raw_response = search_query(client, "a", "{}", "50") time = raw_response[0]["time"] _id = raw_response[0]["_id"] assert isinstance(_id, str) assert isinstance(time, str) class TestFormatSort: def test_format_sort_correctly(self): """ Given: a sort string in the correct format Then: Format the string in the correct format to be used in `pymongo.sort()` """ assert format_sort("field1:asc,field2:desc") == [("field1", 1), ("field2", -1)] assert format_sort("field1:asc") == [("field1", 1)] def test_format_sort_raises_error(self): """ Given: a sort string in the wrong format Then: raise a ValueError """ with pytest.raises(ValueError): format_sort("Wrong:Type") with pytest.raises(ValueError): format_sort("WrongType") def test_pipeline_query_command(mocker): """ Given: collection - where to search. pipeline - json pipeline query When: calling `pipeline_query_command` Then: validate the readable output and context """ client = Client(["aaaaa"], "a", "b", "d") return_value = [ {"title": "test_title", "color": "red", "year": "2019", "_id": "6034a5a62f605638740dba55"}, {"title": "test_title", "color": "yellow", "year": "2020", "_id": "6034a5c52f605638740dba57"}, ] mocker.patch.object(client, "pipeline_query", return_value=return_value) readable_outputs, outputs, raw_response = pipeline_query_command( client=client, collection="test_collection", pipeline='[{"$match": {"title": "test_title"}}]' ) expected_context = [] for item in copy.deepcopy(raw_response): item.update({"collection": "test_collection"}) expected_context.append(item) assert "Total of 2 entries were found in MongoDB collection" in readable_outputs assert outputs.get("MongoDB.Entry(val._id === obj._id && obj.collection === val.collection)") == expected_context class MockResponse: """Mock response for TestUpdateQueryCommands and TestBulkUpdateQueryCommands classes. represents a partial SDK response of the update_entry and bulk_update_entries functions. """ def __init__(self, acknowledged, modified_count, upserted_count, upserted_id=False): self.acknowledged = acknowledged self.modified_count = modified_count self.upserted_count = upserted_count self.upserted_id = upserted_id class TestUpdateQueryCommands: """Class for update_query_command UTs.""" client = Client(["aaaaa"], "a", "b", "d") case_upsert_with_no_matching_entry = ( '{"Name": "dummy"}', '{"$set":{"test":0}}', True, True, MockResponse(True, 0, 0, 1), "A new entry was inserted to the collection.", ) case_upsert_with_one_matching_entry = ( '{"Name": "dummy"}', '{"$set":{"test":0}}', True, True, MockResponse(True, 1, 0, 0), "MongoDB: Total of 1 entries has been modified.", ) case_upsert_with_many_matching_entry = ( '{"Name": "dummy"}', '{"$set":{"test":0}}', False, True, MockResponse(True, 5, 0, 0), "MongoDB: Total of 5 entries has been modified.", ) case_upsert_with_matching_entry_no_modifications = ( '{"Name": "dummy"}', '{"$set":{"test":0}}', True, True, MockResponse(True, 0, 0, 0), "MongoDB: Total of 0 entries has been modified.", ) case_upsert_with_many_matching_entries_update_only_one = ( '{"Name": "dummy"}', '{"$set":{"test":0}}', True, True, MockResponse(True, 1, 0, 0), "MongoDB: Total of 1 entries has been modified.", ) case_no_upsert_with_no_matching_entry = ( '{"Name": "dummy"}', '{"$set":{"test":0}}', True, False, MockResponse(True, 0, 0, 0), "MongoDB: Total of 0 entries has been modified.", ) update_query_cases = [ case_upsert_with_no_matching_entry, case_upsert_with_one_matching_entry, case_upsert_with_many_matching_entry, case_upsert_with_matching_entry_no_modifications, case_upsert_with_many_matching_entries_update_only_one, case_no_upsert_with_no_matching_entry, ] @pytest.mark.parametrize("filter, update, update_one, upsert, response, expected", update_query_cases) def test_update_entry_command(self, mocker, filter, update, update_one, upsert, response, expected, client=client): """ Given: valid arguments When: running mongodb-update command in XSOAR Then: the expected human readable is returned """ mocker.patch.object(client, "update_entry", return_value=response) return_value = update_entry_command( client, "test_collection", filter=filter, update=update, update_one=update_one, upsert=upsert ) assert return_value[0] == expected case_invalid_filter_argument = ( '"Name": "dummy"}', '{"$set":{"test":0}}', MockResponse(True, 0, 0, 0), 'The `filter` argument is not a valid json. Valid input example: `{"key": "value"}`', ) case_invalid_update_argument = ( '{"Name": "dummy"}', '"$set":{"test":0}}', MockResponse(True, 0, 0, 0), 'The `update` argument is not a valid json. Valid input example: `{"$set": {"key": "value"}`', ) case_invalid_response = ( '{"Name": "dummy"}', '{"$set":{"test":0}}', None, "Error occurred when trying to enter update entries.", ) invalid_cases = [case_invalid_filter_argument, case_invalid_update_argument, case_invalid_response] @pytest.mark.parametrize("filter, update, response, expected", invalid_cases) def test_update_entry_command_fail(self, mocker, filter, update, response, expected, client=client): """ Given: invalid arguments When: running mongodb-update command in XSOAR Then: the expected error message is raised """ mocker.patch.object(client, "update_entry", return_value=response) try: update_entry_command(client, "test_collection", filter=filter, update=update) except DemistoException as e: assert str(e) == expected class TestBulkUpdateQueryCommands: """Class for bulk_update_query_command UTs.""" client = Client(["aaaaa"], "a", "b", "d") # valid command arguments case_single_update_args = ('[{"Name": "dummy"}]', '[{"$set":{"test":0}}]', ([{"Name": "dummy"}], [{"$set": {"test": 0}}])) case_simple_bulk_update_args = ( '[{"Name": "dummy1"},{"Name": "dummy2"}]', '[{"$set":{"test":1}},{"$set":{"test":2}}]', ([{"Name": "dummy1"}, {"Name": "dummy2"}], [{"$set": {"test": 1}}, {"$set": {"test": 2}}]), ) case_bulk_update_complex_filter = ( '[{"$and": [{"value":0,"another_value":1}],"Name":"dummy1",\ "less_than": {"$lt":3000}},{"Name":"dummy2"}]', '[{"$set":{"test":1}},{"$set":{"test":2}}]', ( [{"$and": [{"value": 0, "another_value": 1}], "less_than": {"$lt": 3000}, "Name": "dummy1"}, {"Name": "dummy2"}], [{"$set": {"test": 1}}, {"$set": {"test": 2}}], ), ) case_bulk_update_complex_update = ( '[{"Name":"dummy1"},{"Name":"dummy2"}]', '[{"$set":{"test":1,"value":2,"another_value":{"sub_value": 4}}},{"$set":{"test":2}}]', ( [{"Name": "dummy1"}, {"Name": "dummy2"}], [{"$set": {"test": 1, "value": 2, "another_value": {"sub_value": 4}}}, {"$set": {"test": 2}}], ), ) case_bulk_update_context_args = ( [{"Name": "dummy1"}, {"Name": "dummy2"}], [{"$set": {"test": 1}}, {"$set": {"test": 2}}], ([{"Name": "dummy1"}, {"Name": "dummy2"}], [{"$set": {"test": 1}}, {"$set": {"test": 2}}]), ) # invalid command arguments case_missing_brackets = ( '{"Name": "dummy1"},{"Name": "dummy2"}]', '[{"$set":{"test":1}},{"$set":{"test":2}}]', "The `filter` argument must be a json array.", ) case_not_matching_number_of_filters_and_updates = ( '[{"Name": "dummy1"},{"Name": "dummy2"}]', '[{"$set":{"test":1}}]', "The `filter` and `update` arguments must contain the same number of elements.", ) case_invalid_json = ( '{"Name": "dummy1"},{"Name": "dummy2"]', '[{"$set":{"test":1}},{"$set":{"test":2}}]', "The `filter` argument contains an invalid json.", ) @pytest.mark.parametrize( "filter, update, expected_output", [ case_single_update_args, case_simple_bulk_update_args, case_bulk_update_complex_filter, case_bulk_update_complex_update, case_bulk_update_context_args, ], ) def test_parse_and_validate_bulk_update_arguments(self, filter, update, expected_output): """ Given: valid arguments for bulk update command When: running mongodb-bulk-update command in XSOAR Then: parse_and_validate_bulk_update_arguments will parse validate the filter and update arguments """ filter_list, update_list = parse_and_validate_bulk_update_arguments(filter, update) assert filter_list == expected_output[0] assert update_list == expected_output[1] @pytest.mark.parametrize( "filter, update, error_message", [case_missing_brackets, case_not_matching_number_of_filters_and_updates, case_invalid_json], ) def test_parse_and_validate_bulk_update_arguments_fail(self, filter, update, error_message): """ Given: invalid arguments for bulk update command When: running mongodb-bulk-update command in XSOAR Then: parse_and_validate_bulk_update_arguments will raise an error """ with pytest.raises(DemistoException) as e: parse_and_validate_bulk_update_arguments(filter, update) assert error_message in str(e.value) def test_bulk_update_command(self, mocker, client=client, case_simple_bulk_update_args=case_simple_bulk_update_args): """ Given: valid arguments for bulk update command When: running mongodb-bulk-update command in XSOAR Then: the expected human readable is returned """ response = MockResponse(acknowledged=True, modified_count=1, upserted_count=1) mocker.patch.object(client, "bulk_update_entries", return_value=response) return_value = bulk_update_command( client, "test_collection", filter=case_simple_bulk_update_args[0], update=case_simple_bulk_update_args[1] ) excepted_output = "MongoDB: Total of 1 entries has been modified.\ \nMongoDB: Total of 1 entries has been inserted." # 'replace' method is used due to inconsistent spaces in the output assert return_value[0].replace(" ", "") == excepted_output.replace(" ", "") def test_client_initialization_success(): with patch("MongoDB.MongoClient") as mock_mongo_client: mock_db = MagicMock() mock_mongo_client.return_value.get_database.return_value = mock_db client = Client( urls=["mongodb://localhost:27017"], username="testuser", password="testpass", database="testdb", ssl=True, insecure=True, auth_source="admin", timeout=3000, ) assert client.db == mock_db mock_mongo_client.assert_called_once_with( host=["mongodb://localhost:27017"], username="testuser", password="testpass", ssl=True, socketTimeoutMS=3000, tlsAllowInvalidCertificates=True, authSource="admin", ) def test_client_initialization_insecure_without_ssl(): with pytest.raises(DemistoException) as e: Client( urls=["mongodb://localhost:27017"], username="testuser", password="testpass", database="testdb", ssl=False, insecure=True, ) assert e.value.args[0] == '"Trust any certificate (not secure)" must be ticked with "Use TLS/SSL secured connection"' def test_client_initialization_without_auth_source(): with patch("MongoDB.MongoClient") as mock_mongo_client: Client(urls=["mongodb://localhost:27017"], username="testuser", password="testpass", database="testdb", ssl=True) mock_mongo_client.assert_called_once() assert "authSource" not in mock_mongo_client.call_args[1]