# Copyright (c) Microsoft Corporation. # Licensed under the MIT License import json import time import unittest import uuid from unittest import mock from uuid import uuid4 from azure.storage.queue import QueueMessage, QueueClient from azure.kusto.ingest import QueuedIngestClient, IngestionResult, IngestionStatus from azure.kusto.ingest._resource_manager import _ResourceUri from azure.kusto.ingest.status import KustoIngestStatusQueues, SuccessMessage, FailureMessage SAS = "sp=rl&st=2020-05-20T13:38:37Z&se=2020-05-21T13:38:37Z&sv=2019-10-10&sr=c&sig=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" ENDPOINT_SUFFIX = "core.windows.net" def mock_message(success): m = QueueMessage() m.id = uuid4() m.insertion_time = time.time() m.expiration_time = None m.dequeue_count = None if success: content = { "OperationId": str(m.id), "Database": "db1", "Table": "table1", "IngestionSourceId": str(m.id), "IngestionSourcePath": "blob/path", "RootActivityId": "1", "SucceededOn": time.time(), } else: content = { "OperationId": str(m.id), "Database": "db1", "Table": "table1", "IngestionSourceId": str(m.id), "IngestionSourcePath": "blob/path", "RootActivityId": "1", "FailedOn": time.time(), "Details": "", "ErrorCode": "1", "FailureStatus": "", "OriginatesFromUpdatePolicy": "", "ShouldRetry": False, } m.content = json.dumps(content) m.pop_receipt = None m.time_next_visible = None return m def fake_peek_factory(f): def fake_peek(self, max_messages): return f(self.queue_name, max_messages) return fake_peek def fake_receive_factory(f): def fake_receive(self, messages_per_page, *args, **kwargs): return f(self.queue_name, messages_per_page) return fake_receive def fake_delete_factory(f): def fake_delete(self, n): return f(self.queue_name, n) return fake_delete def get_resource_uri(account: str, container: str) -> _ResourceUri: return _ResourceUri("https://{0}.{1}.{2}/{3}?{4}".format(account, "queue", ENDPOINT_SUFFIX, container, SAS)) class StatusQTests(unittest.TestCase): def test_init(self): client = QueuedIngestClient("some-cluster") qs = KustoIngestStatusQueues(client) assert qs.success.message_cls == SuccessMessage assert qs.failure.message_cls == FailureMessage def test_isempty(self): client = QueuedIngestClient("some-cluster") fake_peek = fake_peek_factory( lambda queue_name, num_messages=1: [mock_message(success=True) for _ in range(0, num_messages)] if "qs" in queue_name else [] ) with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues") as mocked_get_success_qs, mock.patch.object( client._resource_manager, "get_failed_ingestions_queues" ) as mocked_get_failed_qs, mock.patch.object(QueueClient, "peek_messages", autospec=True, side_effect=fake_peek) as q_mock: fake_failed_queue = get_resource_uri("mocked_storage_account_f1", "mocked_qf_name") fake_success_queue = get_resource_uri("mocked_storage_account2", "mocked_qs_name") mocked_get_success_qs.return_value = [fake_success_queue] mocked_get_failed_qs.return_value = [fake_failed_queue] qs = KustoIngestStatusQueues(client) assert qs.success.is_empty() is False assert qs.failure.is_empty() is True assert q_mock.call_count == 2 assert q_mock.call_args_list[0][1]["max_messages"] == 2 assert q_mock.call_args_list[1][1]["max_messages"] == 2 def test_peek(self): client = QueuedIngestClient("some-cluster") fake_peek = fake_peek_factory( lambda queue_name, num_messages=1: [ mock_message(success=True) if "qs" in queue_name else mock_message(success=False) for _ in range(0, num_messages) ] ) with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues") as mocked_get_success_qs, mock.patch.object( client._resource_manager, "get_failed_ingestions_queues" ) as mocked_get_failed_qs, mock.patch.object(QueueClient, "peek_messages", autospec=True, side_effect=fake_peek) as q_mock: fake_failed_queue1 = get_resource_uri("mocked_storage_account_f1", "mocked_qf_name") fake_failed_queue2 = get_resource_uri("mocked_storage_account_f2", "mocked_qf_2_name") fake_success_queue = get_resource_uri("mocked_storage_account2", "mocked_qs_name") mocked_get_success_qs.return_value = [fake_success_queue] mocked_get_failed_qs.return_value = [fake_failed_queue1, fake_failed_queue2] qs = KustoIngestStatusQueues(client) peek_success_actual = qs.success.peek() peek_failure_actual = qs.failure.peek(6) assert len(peek_success_actual) == 1 for m in peek_failure_actual: assert isinstance(m, FailureMessage) is True for m in peek_success_actual: assert isinstance(m, SuccessMessage) is True assert len(peek_failure_actual) == 6 actual = {} assert len(QueueClient.peek_messages.call_args_list) == 3 for call_args in q_mock.call_args_list: actual[call_args[0][0].queue_name] = actual.get(call_args[0][0].queue_name, 0) + call_args[1]["max_messages"] assert actual[fake_failed_queue2.object_name] == 4 assert actual[fake_failed_queue1.object_name] == 4 assert actual[fake_success_queue.object_name] == 2 def test_pop(self): client = QueuedIngestClient("some-cluster") fake_receive = fake_receive_factory( lambda queue_name, num_messages=1: [ mock_message(success=True) if "qs" in queue_name else mock_message(success=False) for _ in range(0, num_messages) ] ) with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues") as mocked_get_success_qs, mock.patch.object( client._resource_manager, "get_failed_ingestions_queues" ) as mocked_get_failed_qs, mock.patch.object( QueueClient, "receive_messages", autospec=True, side_effect=fake_receive, ) as q_receive_mock, mock.patch.object( QueueClient, "delete_message", return_value=None ) as q_del_mock: fake_failed_queue1 = get_resource_uri("mocked_storage_account_f1", "mocked_qf_name") fake_failed_queue2 = get_resource_uri("mocked_storage_account_f2", "mocked_qf_2_name") fake_success_queue = get_resource_uri("mocked_storage_account2", "mocked_qs_name") mocked_get_success_qs.return_value = [fake_success_queue] mocked_get_failed_qs.return_value = [fake_failed_queue1, fake_failed_queue2] qs = KustoIngestStatusQueues(client) get_success_actual = qs.success.pop() get_failure_actual = qs.failure.pop(6) assert len(get_success_actual) == 1 assert len(get_failure_actual) == 6 for m in get_failure_actual: assert isinstance(m, FailureMessage) for m in get_success_actual: assert isinstance(m, SuccessMessage) assert q_receive_mock.call_count == 3 assert q_del_mock.call_count == len(get_success_actual) + len(get_failure_actual) assert q_receive_mock.call_args_list[0][1]["messages_per_page"] == 2 actual = { q_receive_mock.call_args_list[1][0][0].queue_name: q_receive_mock.call_args_list[1][1]["messages_per_page"], q_receive_mock.call_args_list[2][0][0].queue_name: q_receive_mock.call_args_list[2][1]["messages_per_page"], } assert actual[fake_failed_queue2.object_name] == 4 assert actual[fake_failed_queue1.object_name] == 4 client.close() def test_pop_unbalanced_queues(self): client = QueuedIngestClient("some-cluster") fake_receive = fake_receive_factory( lambda queue_name, messages_per_page=1: [mock_message(success=False) for _ in range(0, messages_per_page)] if "1" in queue_name else [] ) with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues"), mock.patch.object( client._resource_manager, "get_failed_ingestions_queues" ) as mocked_get_failed_qs, mock.patch.object( QueueClient, "receive_messages", autospec=True, side_effect=fake_receive, ) as q_receive_mock, mock.patch.object( QueueClient, "delete_message", return_value=None ): fake_failed_queue1 = get_resource_uri("mocked_storage_account_f1", "mocked_qf_1_name") fake_failed_queue2 = get_resource_uri("mocked_storage_account_f2", "mocked_qf_2_name") mocked_get_failed_qs.return_value = [fake_failed_queue1, fake_failed_queue2] qs = KustoIngestStatusQueues(client) get_failure_actual = qs.failure.pop(6) assert len(get_failure_actual) == 6 for m in get_failure_actual: assert isinstance(m, FailureMessage) assert q_receive_mock.call_count == 3 actual = {} for call_args in q_receive_mock.call_args_list: actual[call_args[0][0].queue_name] = actual.get(call_args[0][0].queue_name, 0) + call_args[1]["messages_per_page"] assert actual[fake_failed_queue2.object_name] + actual[fake_failed_queue1.object_name] == (4 + 4 + 6) def test_status_sanity(self): i = IngestionResult(IngestionStatus.QUEUED, "db", "table", uuid.UUID("11111111-1111-1111-1111-111111111111"), "path") assert ( repr(i) == "IngestionResult(status=IngestionStatus.QUEUED, database=db, table=table, source_id=11111111-1111-1111-1111-111111111111, obfuscated_blob_uri=path)" ) no_path = IngestionResult(IngestionStatus.QUEUED, "db", "table", uuid.UUID("11111111-1111-1111-1111-111111111111")) assert repr(no_path) == "IngestionResult(status=IngestionStatus.QUEUED, database=db, table=table, source_id=11111111-1111-1111-1111-111111111111)"