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 = _ResourceUri(
                "mocked_storage_account_f1",
                OBJECT_TYPE,
                "queue",
                "mocked_qf_1_name",
                ENDPOINT_SUFFIX,
            )
            fake_failed_queue2 = _ResourceUri(
                "mocked_storage_account_f2",
                OBJECT_TYPE,
                "queue",
                "mocked_qf_2_name",
                ENDPOINT_SUFFIX,
            )

            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_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 = _ResourceUri(
                "mocked_storage_account1",
                OBJECT_TYPE,
                "queue",
                "mocked_qf_name",
                ENDPOINT_SUFFIX,
            )
            fake_success_queue = _ResourceUri(
                "mocked_storage_account2",
                OBJECT_TYPE,
                "queue",
                "mocked_qs_name",
                ENDPOINT_SUFFIX,
            )

            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_isempty(self, mocked_q_peek_messages):
        client = KustoIngestClient("some-cluster")
        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:

            fake_failed_queue = _ResourceUri("mocked_storage_account1",
                                             "queue", "mocked_qf_name",
                                             "mocked_sas")
            fake_success_queue = _ResourceUri("mocked_storage_account2",
                                              "queue", "mocked_qs_name",
                                              "mocked_sas")

            mocked_get_success_qs.return_value = [fake_success_queue]
            mocked_get_failed_qs.return_value = [fake_failed_queue]

            mocked_q_peek_messages.side_effect = (
                lambda queue_name, num_messages=1: []
                if queue_name == fake_failed_queue.object_name else
                [QueueMessage() for _ in range(0, num_messages)])

            qs = KustoIngestStatusQueues(client)

            assert qs.success.is_empty() == False
            assert qs.failure.is_empty() == True

            assert mocked_q_peek_messages.call_count == 2
            assert mocked_q_peek_messages.call_args_list[0][0][
                0] == fake_success_queue.object_name
            assert mocked_q_peek_messages.call_args_list[0][1][
                "num_messages"] == 2

            assert mocked_q_peek_messages.call_args_list[1][0][
                0] == fake_failed_queue.object_name
            assert mocked_q_peek_messages.call_args_list[1][1][
                "num_messages"] == 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 = _ResourceUri(
                "mocked_storage_account_f1",
                OBJECT_TYPE,
                "queue",
                "mocked_qf_1_name",
                ENDPOINT_SUFFIX,
            )
            fake_failed_queue2 = _ResourceUri(
                "mocked_storage_account_f2",
                OBJECT_TYPE,
                "queue",
                "mocked_qf_2_name",
                ENDPOINT_SUFFIX,
            )
            fake_success_queue = _ResourceUri(
                "mocked_storage_account2",
                OBJECT_TYPE,
                "queue",
                "mocked_qs_name",
                ENDPOINT_SUFFIX,
            )

            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
    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 = _ResourceUri(
                "mocked_storage_account_f1",
                OBJECT_TYPE,
                "queue",
                "mocked_qf_1_name",
                ENDPOINT_SUFFIX,
            )
            fake_failed_queue2 = _ResourceUri(
                "mocked_storage_account_f2",
                OBJECT_TYPE,
                "queue",
                "mocked_qf_2_name",
                ENDPOINT_SUFFIX,
            )
            fake_success_queue = _ResourceUri(
                "mocked_storage_account2",
                OBJECT_TYPE,
                "queue",
                "mocked_qs_name",
                ENDPOINT_SUFFIX,
            )

            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_peek(self, mocked_q_peek_messages):
        client = KustoIngestClient("some-cluster")
        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:

            fake_failed_queue1 = _ResourceUri("mocked_storage_account_f1",
                                              "queue", "mocked_qf_1_name",
                                              "mocked_sas")
            fake_failed_queue2 = _ResourceUri("mocked_storage_account_f2",
                                              "queue", "mocked_qf_2_name",
                                              "mocked_sas")
            fake_success_queue = _ResourceUri("mocked_storage_account2",
                                              "queue", "mocked_qs_name",
                                              "mocked_sas")

            mocked_get_success_qs.return_value = [fake_success_queue]
            mocked_get_failed_qs.return_value = [
                fake_failed_queue1, fake_failed_queue2
            ]

            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 = str(
                    base64.b64encode(
                        json.dumps(content).encode("utf-8")).decode("utf-8"))
                m.pop_receipt = None
                m.time_next_visible = None

                return m

            mocked_q_peek_messages.side_effect = lambda queue_name, num_messages=1: [
                mock_message(success=True)
                if queue_name in [fake_success_queue.object_name] else
                mock_message(success=False) for i in range(0, num_messages)
            ]

            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) == True

            for m in peek_success_actual:
                assert isinstance(m, SuccessMessage) == True

            assert len(peek_failure_actual) == 6

            actual = {}

            assert len(mocked_q_peek_messages.call_args_list) == 3

            for call_args in mocked_q_peek_messages.call_args_list:
                actual[call_args[0][0]] = actual.get(
                    call_args[0][0], 0) + call_args[1]["num_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