def test_retrieve_data_from_xdr_sanity_check(mocker):
"""
Given:
All the required args.
When:
Executing retrieve_data_from_xdr function (first case of executing query).
Then:
Ensure executeCommand was called once with the correct query and time_frame args.
"""
# Mock input arguments
args = {"time_frame": "7 days", "indicator": "1.2.3.4", "query_name": "Test Query", "data_set": "xdr_data"}
# Mock return values
initial_response = [
{"Type": 1, "Contents": {"status": "PENDING"}, "Metadata": {"pollingArgs": {"query_id": "abc123", "query_name": "TEST"}}}
]
completed_response = [{"Type": 1, "Contents": {"status": "COMPLETED"}, "HumanReadable": "Query results here"}]
# Patching dependencies
mock_execute = mocker.patch(
"SearchIndicatorInEvents.demisto.executeCommand", side_effect=[initial_response, completed_response]
)
from SearchIndicatorInEvents import retrieve_data_from_xdr
poll_result = retrieve_data_from_xdr(args)
assert mock_execute.call_count == 1
call_args = mock_execute.mock_calls[0][2]["args"]
assert call_args["query"] == 'search "1.2.3.4" dataset = xdr_data'
assert call_args["time_frame"] == "7 days"
assert poll_result.scheduled_command._args["query_id"] == "abc123"
def test_check_status(mocker):
"""
Given:
- All the required args
When:
- Executing retrieve_data_from_xdr function (second case of checking status)
Then:
- Ensure executeCommand was called once
- Ensure the correct args were sent in executeCommand call, query_id and the right command
- Ensure the readable_output is correct
"""
# Mock input arguments
args = {
"time_frame": "7 days",
"indicator": "1.2.3.4",
"query_name": "Test Query",
"data_set": "xdr_data",
"query_id": "abc123",
}
# Mock return values
completed_response = [{"Type": 1, "Contents": {"status": "COMPLETED", "results": []}}]
# Patching dependencies
mock_execute = mocker.patch("SearchIndicatorInEvents.demisto.executeCommand", side_effect=[completed_response])
from SearchIndicatorInEvents import retrieve_data_from_xdr
poll_result = retrieve_data_from_xdr(args)
assert mock_execute.call_count == 1
assert mock_execute.mock_calls[0][2]["args"] == {"query_id": "abc123"}
assert mock_execute.mock_calls[0][2]["command"] == "xdr-xql-get-query-results"
assert poll_result.readable_output == "job ID abc123 is finished!"