#
# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
# or more contributor license agreements. Licensed under the Elastic License 2.0;
# you may not use this file except in compliance with the Elastic License 2.0.
#
import asyncio
import datetime
import itertools
import json
import uuid
from copy import deepcopy
from decimal import Decimal
from unittest import mock
from unittest.mock import ANY, AsyncMock, Mock, call, patch

import pytest
from elastic_transport import JsonSerializer
from elasticsearch import ApiError, BadRequestError

from connectors.es.management_client import ESManagementClient
from connectors.es.sink import (
    BIN_DOCS_DOWNLOADED,
    BULK_OPERATIONS,
    BULK_RESPONSES,
    CREATES_QUEUED,
    DELETES_QUEUED,
    DOCS_DROPPED_TOO_LARGE,
    DOCS_EXTRACTED,
    END_DOCS,
    OP_DELETE,
    OP_INDEX,
    OP_UPDATE,
    RESULT_SUCCESS,
    UPDATES_QUEUED,
    AsyncBulkRunningError,
    ElasticsearchOverloadedError,
    Extractor,
    ForceCanceledError,
    Sink,
    SyncOrchestrator,
)
from connectors.protocol import JobType, Pipeline
from connectors.protocol.connectors import (
    DELETED_DOCUMENT_COUNT,
    INDEXED_DOCUMENT_COUNT,
    INDEXED_DOCUMENT_VOLUME,
)
from connectors.utils import ErrorMonitor, TooManyErrors
from tests.commons import AsyncIterator

INDEX = "some-index"
TIMESTAMP = datetime.datetime(year=2023, month=1, day=1)
NO_FILTERING = ()
DOC_ONE_ID = 1
DOC_TWO_ID = 2
DOC_THREE_ID = 3

DOC_ONE = {"_id": DOC_ONE_ID, "_timestamp": TIMESTAMP}

DOC_ONE_DIFFERENT_TIMESTAMP = {
    "_id": DOC_ONE_ID,
    "_timestamp": TIMESTAMP + datetime.timedelta(days=1),
}

DOC_TWO = {"_id": 2, "_timestamp": TIMESTAMP}
DOC_THREE = {"_id": 3, "_timestamp": TIMESTAMP}
DOC_FOUR = {"_id": 4, "_timestamp": TIMESTAMP}

BULK_ACTION_ERROR = "some error"

SYNC_RULES_ENABLED = True
SYNC_RULES_DISABLED = False
CONTENT_EXTRACTION_ENABLED = True
CONTENT_EXTRACTION_DISABLED = False


def failed_bulk_action(doc_id, action, result, error=BULK_ACTION_ERROR):
    return {action: {"_id": doc_id, "result": result, "error": error}}


def successful_bulk_action(doc_id, action, result):
    return {action: {"_id": doc_id, "result": result}}


def successful_action_log_message(doc_id, action, result):
    return f"Successfully executed '{action}' on document with id '{doc_id}'. Result: {result}"


def successful_operation_with_non_successful_result_log_message(doc_id, action, result):
    return f"Executed '{action}' on document with id '{doc_id}', but got non-successful result: {result}"


def failed_action_log_message(doc_id, action, result, error=BULK_ACTION_ERROR):
    return (
        f"Failed to execute '{action}' on document with id '{doc_id}'. Error: {error}"
    )


@patch("connectors.es.sink.CANCELATION_TIMEOUT", -1)
@pytest.mark.asyncio
async def test_prepare_content_index_raise_error_when_index_creation_failed(
    mock_responses,
):
    index_name = "search-new-index"
    config = {"host": "http://nowhere.com:9200", "user": "tarek", "password": "blah"}
    headers = {"X-Elastic-Product": "Elasticsearch"}
    mock_responses.post(
        "http://nowhere.com:9200/.elastic-connectors/_refresh", headers=headers
    )

    # not found
    mock_responses.get(
        f"http://nowhere.com:9200/{index_name}?ignore_unavailable=true",
        headers=headers,
        status=200,
        body="{}",
    )

    mock_responses.put(
        f"http://nowhere.com:9200/{index_name}",
        payload={"_id": "1"},
        headers=headers,
        status=400,
    )

    es = SyncOrchestrator(config)
    es._sink = Mock()
    es._extractor = Mock()

    with pytest.raises(BadRequestError):
        await es.prepare_content_index(index_name)
    await es.close()


@patch("connectors.es.sink.CANCELATION_TIMEOUT", -1)
@pytest.mark.asyncio
async def test_prepare_content_index_create_index(
    mock_responses,
):
    index_name = "search-new-index"
    language_code = "jp"
    config = {"host": "http://nowhere.com:9200", "user": "tarek", "password": "blah"}
    headers = {"X-Elastic-Product": "Elasticsearch"}
    mock_responses.post(
        "http://nowhere.com:9200/.elastic-connectors/_refresh", headers=headers
    )

    # not found
    mock_responses.get(
        f"http://nowhere.com:9200/{index_name}?ignore_unavailable=true",
        headers=headers,
        status=200,
        body="{}",
    )

    mock_responses.put(
        f"http://nowhere.com:9200/{index_name}",
        payload={"_id": "1"},
        headers=headers,
    )

    mock_responses.put(
        f"http://nowhere.com:9200/{index_name}/_settings",
        payload={"_id": "1"},
        headers=headers,
    )

    es = SyncOrchestrator(config)
    es._sink = Mock()
    es._extractor = Mock()

    create_index_result = asyncio.Future()
    create_index_result.set_result({"acknowledged": True})

    with mock.patch.object(
        es.es_management_client,
        "create_content_index",
        return_value=create_index_result,
    ) as create_index_mock:
        await es.prepare_content_index(index_name, language_code)

        await es.close()

        create_index_mock.assert_called_with(index_name, language_code)


@patch("connectors.es.sink.CANCELATION_TIMEOUT", -1)
@pytest.mark.asyncio
async def test_prepare_content_index(mock_responses):
    language_code = "en"
    config = {"host": "http://nowhere.com:9200", "user": "tarek", "password": "blah"}
    headers = {"X-Elastic-Product": "Elasticsearch"}
    index_name = "search-new-index"

    response = {
        index_name: {
            "mappings": {},
            "settings": {},
        }
    }
    mock_responses.get(
        f"http://nowhere.com:9200/{index_name}?ignore_unavailable=true",
        headers=headers,
        status=200,
        body=json.dumps(response),
    )

    es = SyncOrchestrator(config)
    es._sink = Mock()
    es._extractor = Mock()
    with mock.patch.object(
        es.es_management_client,
        "create_content_index",
    ) as create_index_mock:
        await es.prepare_content_index(index_name, language_code)

        await es.close()

        create_index_mock.assert_not_called()


def set_responses(mock_responses, ts=None):
    if ts is None:
        ts = datetime.datetime.now().isoformat()
    headers = {"X-Elastic-Product": "Elasticsearch"}
    mock_responses.head(
        "http://nowhere.com:9200/search-some-index",
        headers=headers,
    )

    mock_responses.get(
        "http://nowhere.com:9200/search-some-index",
        payload={"_id": "1"},
        headers=headers,
    )
    source_1 = {"id": "1", "_timestamp": ts}
    source_2 = {"id": "2", "_timestamp": ts}

    mock_responses.post(
        "http://nowhere.com:9200/search-some-index/_search?scroll=5m",
        payload={
            "_scroll_id": "1234",
            "_shards": {},
            "hits": {
                "hits": [
                    {"_id": "1", "_source": source_1},
                    {"_id": "2", "_source": source_2},
                ]
            },
        },
        headers=headers,
    )
    mock_responses.delete(
        "http://nowhere.com:9200/_search/scroll",
        headers=headers,
    )
    mock_responses.post(
        "http://nowhere.com:9200/_search/scroll",
        payload={
            "hits": {
                "hits": [
                    {"_id": "1", "_source": source_1},
                    {"_id": "2", "_source": source_2},
                ]
            }
        },
        headers=headers,
    )

    mock_responses.put(
        "http://nowhere.com:9200/_bulk?pipeline=search-default-ingestion",
        payload={
            "took": 7,
            "errors": False,
            "items": [
                {
                    "index": {
                        "_index": "test",
                        "_id": "1",
                        "_version": 1,
                        "result": "created",
                        "forced_refresh": False,
                        "status": 200,
                    }
                },
                {
                    "delete": {
                        "_index": "test",
                        "_id": "3",
                        "result": "deleted",
                        "status": 200,
                    }
                },
                {
                    "update": {
                        "_index": "test",
                        "_id": "3",
                        "result": "updated",
                        "status": 200,
                    }
                },
            ],
        },
        headers=headers,
    )


@patch("connectors.es.sink.CANCELATION_TIMEOUT", -1)
@pytest.mark.asyncio
async def test_async_bulk(mock_responses):
    config = {"host": "http://nowhere.com:9200", "user": "tarek", "password": "blah"}
    set_responses(mock_responses)

    es = SyncOrchestrator(config)
    pipeline = Pipeline({})

    async def get_docs():
        async def _dl_none(doit=True, timestamp=None):
            return None

        async def _dl(doit=True, timestamp=None):
            if not doit:
                return
            return {"TEXT": "DATA", "_timestamp": timestamp, "_id": "1"}

        yield (
            {
                "_id": "1",
                "_timestamp": datetime.datetime.now().isoformat(),
            },
            _dl,
            "index",
        )
        yield {"_id": "3"}, _dl_none, "index"

    await es.async_bulk("search-some-index", get_docs(), pipeline, JobType.FULL)
    while not es.done():
        await asyncio.sleep(0.1)

    ingestion_stats = es.ingestion_stats()

    assert ingestion_stats == {
        CREATES_QUEUED: 1,
        BIN_DOCS_DOWNLOADED: 1,
        UPDATES_QUEUED: 1,
        DELETES_QUEUED: 1,
        f"{BULK_OPERATIONS}.{OP_INDEX}": 2,
        f"{BULK_OPERATIONS}.{OP_DELETE}": 1,
        INDEXED_DOCUMENT_COUNT: 2,
        INDEXED_DOCUMENT_VOLUME: ANY,
        DELETED_DOCUMENT_COUNT: 1,
        f"{BULK_RESPONSES}.{OP_DELETE}": 1,
        f"{BULK_RESPONSES}.{OP_INDEX}": 1,
        f"{BULK_RESPONSES}.{OP_UPDATE}": 1,
        DOCS_EXTRACTED: 2,
        RESULT_SUCCESS: 3,
    }

    # 2nd sync
    set_responses(mock_responses)
    with pytest.raises(AsyncBulkRunningError):
        await es.async_bulk("search-some-index", get_docs(), pipeline, JobType.FULL)

    await es.close()


def index_operation(doc):
    # deepcopy as get_docs mutates docs
    doc_copy = deepcopy(doc)
    doc_id = str(doc_copy.pop("_id"))
    doc_copy["id"] = doc_id

    return {"_op_type": "index", "_index": INDEX, "_id": doc_id, "doc": doc_copy}


def update_operation(doc):
    # deepcopy as get_docs mutates docs
    doc_copy = deepcopy(doc)
    doc_id = str(doc_copy.pop("_id"))
    doc_copy["id"] = doc_id

    return {"_op_type": "update", "_index": INDEX, "_id": doc_id, "doc": doc_copy}


def delete_operation(doc):
    return {"_op_type": "delete", "_index": INDEX, "_id": str(doc["_id"])}


def end_docs_operation():
    return "END_DOCS"


def created(doc_count):
    """Used for test readability."""
    return doc_count


def updated(doc_count):
    """Used for test readability."""
    return doc_count


def deleted(doc_count):
    """Used for test readability."""
    return doc_count


def total_downloads(count):
    """Used for test readability."""
    return count


async def queue_mock():
    queue = Mock()
    future = asyncio.Future()
    future.set_result(1)
    queue.put = Mock(return_value=future)
    return queue


def lazy_download_fake(doc):
    async def lazy_download(**kwargs):
        # also deepcopy to prevent side effects as docs get mutated in get_docs
        return deepcopy(doc)

    return lazy_download


def crashing_lazy_download_fake():
    async def lazy_download(**kwargs):
        msg = "Could not download"
        raise Exception(msg)

    return lazy_download


def queue_called_with_operations(queue, operations):
    expected_calls = [call(operation) for operation in operations]
    actual_calls = queue.put.call_args_list

    return actual_calls == expected_calls and queue.put.call_count == len(
        expected_calls
    )


async def basic_rule_engine_mock(return_values):
    basic_rule_engine = Mock()
    basic_rule_engine.should_ingest = (
        Mock(side_effect=list(return_values))
        if len(return_values)
        else Mock(return_value=True)
    )
    return basic_rule_engine


async def lazy_downloads_mock():
    lazy_downloads = Mock()
    future = asyncio.Future()
    future.set_result(1)
    lazy_downloads.put = Mock(return_value=future)
    return lazy_downloads


async def setup_extractor(
    queue,
    basic_rule_engine=None,
    sync_rules_enabled=False,
    content_extraction_enabled=False,
):
    config = {
        "username": "elastic",
        "password": "changeme",
        "host": "http://nowhere.com:9200",
    }
    # filtering content doesn't matter as the BasicRuleEngine behavior is mocked
    filter_mock = Mock()
    filter_mock.get_active_filter = Mock(return_value={})
    extractor = Extractor(
        ESManagementClient(config),
        queue,
        INDEX,
        filter_=filter_mock,
        content_extraction_enabled=content_extraction_enabled,
    )
    extractor.basic_rule_engine = basic_rule_engine if sync_rules_enabled else None
    return extractor


@pytest.mark.parametrize(
    "existing_docs, docs_from_source, doc_should_ingest, sync_rules_enabled, content_extraction_enabled, expected_queue_operations, "
    "expected_total_docs_updated, expected_total_docs_created, expected_total_docs_deleted, expected_total_downloads",
    [
        (
            # no docs exist, data source has doc 1 -> ingest doc 1
            [],
            [(DOC_ONE, None, "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
        (
            # doc 1 is present, nothing in source -> delete one doc
            [DOC_ONE],
            [],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [delete_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(0),
            deleted(1),
            total_downloads(0),
        ),
        (
            # doc 2 is present, data source only returns doc 1 -> delete doc 2 and ingest doc 1
            [DOC_TWO],
            [(DOC_ONE, None, "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE), delete_operation(DOC_TWO), end_docs_operation()],
            updated(0),
            created(1),
            deleted(1),
            total_downloads(0),
        ),
        (
            # doc 1 is present, data source also has doc 1 with the same timestamp -> doc one is updated
            [DOC_ONE],
            [(DOC_ONE, None, "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [
                # update happens through overwriting
                index_operation(DOC_ONE),
                end_docs_operation(),
            ],
            updated(1),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            # doc 1 is present, data source has doc 1 with different timestamp -> doc is updated
            [DOC_ONE],
            [(DOC_ONE_DIFFERENT_TIMESTAMP, None, "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [
                # update happens through overwriting
                index_operation(DOC_ONE_DIFFERENT_TIMESTAMP),
                end_docs_operation(),
            ],
            updated(1),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            [],
            [(DOC_ONE, None, "index")],
            # filter out doc 1
            (False,),
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [
                # should not ingest doc 1
                end_docs_operation()
            ],
            updated(0),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            # doc 1 is present, data source has doc 1, doc 1 should be filtered -> delete doc 1
            [DOC_ONE],
            [(DOC_ONE, None, "index")],
            # filter out doc 1
            (False,),
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [
                delete_operation(DOC_ONE),
                end_docs_operation(),
            ],
            updated(0),
            created(0),
            deleted(1),
            total_downloads(0),
        ),
        (
            # doc 1 is present, data source has doc 1 (different timestamp), doc 1 should be filtered -> delete doc 1
            [DOC_ONE],
            [(DOC_ONE_DIFFERENT_TIMESTAMP, None, "index")],
            # still filter out doc 1
            (False,),
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [delete_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(0),
            deleted(1),
            total_downloads(0),
        ),
        (
            [],
            [(DOC_ONE, None, "index")],
            # filtering throws exception
            [Exception()],
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            ["EXTRACTOR_ERROR"],
            updated(0),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            [],
            # no doc present, lazy download for doc 1 specified -> index and increment the downloads counter
            [(DOC_ONE, lazy_download_fake(DOC_ONE), "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(1),
        ),
        (
            # doc 1 present, data source has doc 1 -> lazy download occurs
            [DOC_ONE],
            [(DOC_ONE, lazy_download_fake(DOC_ONE), "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(1),
            created(0),
            deleted(0),
            total_downloads(1),
        ),
        (
            # doc 1 present, data source has doc 1 with different timestamp
            # lazy download coroutine present -> execute lazy download
            [DOC_ONE],
            [
                (
                    DOC_ONE_DIFFERENT_TIMESTAMP,
                    lazy_download_fake(DOC_ONE_DIFFERENT_TIMESTAMP),
                    "index",
                )
            ],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE_DIFFERENT_TIMESTAMP), end_docs_operation()],
            updated(1),
            created(0),
            deleted(0),
            total_downloads(1),
        ),
        (
            [],
            [(DOC_ONE, None, "index")],
            # filter out doc 1
            (False,),
            SYNC_RULES_DISABLED,
            CONTENT_EXTRACTION_ENABLED,
            [
                # should ingest doc 1 as sync rules are disabled
                index_operation(DOC_ONE),
                end_docs_operation(),
            ],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
        (
            # content_extraction_enabled is false and no download is provided,
            # indexing should still work, but nothing should be downloaded
            [],
            [(DOC_ONE, None, "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_DISABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
        (
            # content_extraction_enabled is false but a download is also provided,
            # indexing should still work, but nothing should be downloaded
            [],
            [(DOC_ONE, lazy_download_fake(DOC_ONE), "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_DISABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
    ],
)
@mock.patch(
    "connectors.es.management_client.ESManagementClient.yield_existing_documents_metadata"
)
@pytest.mark.asyncio
async def test_get_docs(
    yield_existing_documents_metadata,
    existing_docs,
    docs_from_source,
    doc_should_ingest,
    sync_rules_enabled,
    content_extraction_enabled,
    expected_queue_operations,
    expected_total_docs_updated,
    expected_total_docs_created,
    expected_total_docs_deleted,
    expected_total_downloads,
):
    lazy_downloads = await lazy_downloads_mock()

    yield_existing_documents_metadata.return_value = AsyncIterator(
        [(str(doc["_id"]), doc["_timestamp"]) for doc in existing_docs]
    )

    with mock.patch("connectors.utils.ConcurrentTasks", return_value=lazy_downloads):
        queue = await queue_mock()
        basic_rule_engine = await basic_rule_engine_mock(doc_should_ingest)

        # deep copying docs is needed as get_docs mutates the document ids which has side effects on other test
        # instances
        doc_generator = AsyncIterator([deepcopy(doc) for doc in docs_from_source])

        extractor = await setup_extractor(
            queue,
            basic_rule_engine=basic_rule_engine,
            sync_rules_enabled=sync_rules_enabled,
            content_extraction_enabled=content_extraction_enabled,
        )

        await extractor.run(doc_generator, JobType.FULL)

        assert extractor.counters.get(UPDATES_QUEUED) == expected_total_docs_updated
        assert extractor.counters.get(CREATES_QUEUED) == expected_total_docs_created
        assert extractor.counters.get(DELETES_QUEUED) == expected_total_docs_deleted
        assert extractor.counters.get(BIN_DOCS_DOWNLOADED) == expected_total_downloads

        assert queue_called_with_operations(queue, expected_queue_operations)


@pytest.mark.parametrize(
    "docs_from_source, doc_should_ingest, sync_rules_enabled, content_extraction_enabled, expected_queue_operations, "
    "expected_total_docs_updated, expected_total_docs_created, expected_total_docs_deleted, expected_total_downloads",
    [
        (
            [],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [end_docs_operation()],
            updated(0),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, None, "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, None, "update")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [update_operation(DOC_ONE), end_docs_operation()],
            updated(1),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, None, "delete")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [delete_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(0),
            deleted(1),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, None, "index"), (DOC_TWO, None, "delete")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE), delete_operation(DOC_TWO), end_docs_operation()],
            updated(0),
            created(1),
            deleted(1),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, None, "index")],
            # filter out doc 1
            (False,),
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [
                # should not ingest doc 1
                end_docs_operation()
            ],
            updated(0),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, None, "index")],
            # filtering throws exception
            [Exception()],
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            ["EXTRACTOR_ERROR"],
            updated(0),
            created(0),
            deleted(0),
            total_downloads(0),
        ),
        (
            # no doc present, lazy download for doc 1 specified -> index and increment the downloads counter
            [(DOC_ONE, lazy_download_fake(DOC_ONE), "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_ENABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(1),
        ),
        (
            [(DOC_ONE, None, "index")],
            # filter out doc 1
            (False,),
            SYNC_RULES_DISABLED,
            CONTENT_EXTRACTION_ENABLED,
            [
                # should ingest doc 1 as sync rules are disabled
                index_operation(DOC_ONE),
                end_docs_operation(),
            ],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, None, "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_DISABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
        (
            [(DOC_ONE, lazy_download_fake(DOC_ONE), "index")],
            NO_FILTERING,
            SYNC_RULES_ENABLED,
            CONTENT_EXTRACTION_DISABLED,
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
            total_downloads(0),
        ),
    ],
)
@pytest.mark.asyncio
async def test_get_docs_incrementally(
    docs_from_source,
    doc_should_ingest,
    sync_rules_enabled,
    content_extraction_enabled,
    expected_queue_operations,
    expected_total_docs_updated,
    expected_total_docs_created,
    expected_total_docs_deleted,
    expected_total_downloads,
):
    lazy_downloads = await lazy_downloads_mock()

    with mock.patch("connectors.utils.ConcurrentTasks", return_value=lazy_downloads):
        queue = await queue_mock()
        basic_rule_engine = await basic_rule_engine_mock(doc_should_ingest)

        # deep copying docs is needed as get_docs mutates the document ids which has side effects on other test
        # instances
        doc_generator = AsyncIterator([deepcopy(doc) for doc in docs_from_source])

        extractor = await setup_extractor(
            queue,
            basic_rule_engine=basic_rule_engine,
            sync_rules_enabled=sync_rules_enabled,
            content_extraction_enabled=content_extraction_enabled,
        )

        await extractor.run(doc_generator, JobType.INCREMENTAL)

        assert extractor.counters.get(UPDATES_QUEUED) == expected_total_docs_updated
        assert extractor.counters.get(CREATES_QUEUED) == expected_total_docs_created
        assert extractor.counters.get(DELETES_QUEUED) == expected_total_docs_deleted
        assert extractor.counters.get(BIN_DOCS_DOWNLOADED) == expected_total_downloads

        assert queue_called_with_operations(queue, expected_queue_operations)


@pytest.mark.parametrize(
    "existing_docs, docs_from_source, expected_queue_operations, "
    "expected_total_docs_updated, expected_total_docs_created, expected_total_docs_deleted",
    [
        (
            [],
            [],
            [end_docs_operation()],
            updated(0),
            created(0),
            deleted(0),
        ),
        (
            [],
            [(DOC_ONE, None, None)],
            [index_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(1),
            deleted(0),
        ),
        (
            [DOC_ONE_DIFFERENT_TIMESTAMP],
            [(DOC_ONE, None, None)],
            [update_operation(DOC_ONE), end_docs_operation()],
            updated(1),
            created(0),
            deleted(0),
        ),
        (
            [DOC_ONE],
            [],
            [delete_operation(DOC_ONE), end_docs_operation()],
            updated(0),
            created(0),
            deleted(1),
        ),
        (
            [DOC_TWO],
            [(DOC_ONE, None, None)],
            [index_operation(DOC_ONE), delete_operation(DOC_TWO), end_docs_operation()],
            updated(0),
            created(1),
            deleted(1),
        ),
        (
            [DOC_ONE_DIFFERENT_TIMESTAMP, DOC_TWO],
            [(DOC_ONE, None, None)],
            [
                update_operation(DOC_ONE),
                delete_operation(DOC_TWO),
                end_docs_operation(),
            ],
            updated(1),
            created(0),
            deleted(1),
        ),
        (
            [DOC_ONE_DIFFERENT_TIMESTAMP, DOC_TWO],
            [(DOC_ONE, None, None), (DOC_THREE, None, None)],
            [
                update_operation(DOC_ONE),
                index_operation(DOC_THREE),
                delete_operation(DOC_TWO),
                end_docs_operation(),
            ],
            updated(1),
            created(1),
            deleted(1),
        ),
    ],
)
@mock.patch(
    "connectors.es.management_client.ESManagementClient.yield_existing_documents_metadata"
)
@pytest.mark.asyncio
async def test_get_access_control_docs(
    yield_existing_documents_metadata,
    existing_docs,
    docs_from_source,
    expected_queue_operations,
    expected_total_docs_updated,
    expected_total_docs_created,
    expected_total_docs_deleted,
):
    yield_existing_documents_metadata.return_value = AsyncIterator(
        [(str(doc["_id"]), doc["_timestamp"]) for doc in existing_docs]
    )

    queue = await queue_mock()
    # deep copying docs is needed as get_access_control_docs mutates the document ids which has side effects on other test
    # instances
    doc_generator = AsyncIterator([deepcopy(doc) for doc in docs_from_source])

    extractor = await setup_extractor(queue=queue)

    await extractor.run(doc_generator, JobType.ACCESS_CONTROL)

    assert extractor.counters.get(UPDATES_QUEUED) == expected_total_docs_updated
    assert extractor.counters.get(CREATES_QUEUED) == expected_total_docs_created
    assert extractor.counters.get(DELETES_QUEUED) == expected_total_docs_deleted

    assert queue_called_with_operations(queue, expected_queue_operations)


STATS = {
    "index": {"1": 1},
    "update": {"2": 1},
    "delete": {"3": 0},
}

INDEX_ITEM = {"index": {"_id": "1", "result": "created"}}
FAILED_INDEX_ITEM = {"index": {"_id": "1"}}
UPDATE_ITEM = {"update": {"_id": "2", "result": "updated"}}
FAILED_UPDATE_ITEM = {"update": {"_id": "2"}}
DELETE_ITEM = {"delete": {"_id": "3", "result": "deleted"}}
FAILED_DELETE_ITEM = {"delete": {"_id": "3"}}


@pytest.mark.parametrize(
    "res, expected_result",
    [
        (
            {"items": [INDEX_ITEM, UPDATE_ITEM, DELETE_ITEM]},
            {
                "indexed_document_count": 2,
                "indexed_document_volume": 2,
                "deleted_document_count": 1,
            },
        ),
        (
            {"items": [FAILED_INDEX_ITEM, UPDATE_ITEM, DELETE_ITEM]},
            {
                "indexed_document_count": 1,
                "indexed_document_volume": 1,
                "deleted_document_count": 1,
            },
        ),
        (
            {"items": [INDEX_ITEM, FAILED_UPDATE_ITEM, DELETE_ITEM]},
            {
                "indexed_document_count": 1,
                "indexed_document_volume": 1,
                "deleted_document_count": 1,
            },
        ),
        (
            {"items": [INDEX_ITEM, UPDATE_ITEM, FAILED_DELETE_ITEM]},
            {
                "indexed_document_count": 2,
                "indexed_document_volume": 2,
                "deleted_document_count": 0,
            },
        ),
        (
            {"items": [INDEX_ITEM, FAILED_UPDATE_ITEM, FAILED_DELETE_ITEM]},
            {
                "indexed_document_count": 1,
                "indexed_document_volume": 1,
                "deleted_document_count": 0,
            },
        ),
        (
            {"items": [FAILED_INDEX_ITEM, UPDATE_ITEM, FAILED_DELETE_ITEM]},
            {
                "indexed_document_count": 1,
                "indexed_document_volume": 1,
                "deleted_document_count": 0,
            },
        ),
        (
            {"items": [FAILED_INDEX_ITEM, FAILED_UPDATE_ITEM, DELETE_ITEM]},
            {
                "indexed_document_count": 0,
                "indexed_document_volume": 0,
                "deleted_document_count": 1,
            },
        ),
        (
            {"items": [FAILED_INDEX_ITEM, FAILED_UPDATE_ITEM, FAILED_DELETE_ITEM]},
            {
                "indexed_document_count": 0,
                "indexed_document_volume": 0,
                "deleted_document_count": 0,
            },
        ),
    ],
)
def test_bulk_populate_stats(res, expected_result):
    sink = Sink(
        client=None,
        queue=None,
        error_monitor=Mock(),
        chunk_size=0,
        pipeline=None,
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )
    sink._populate_stats(deepcopy(STATS), res)

    assert (
        sink.counters.get(INDEXED_DOCUMENT_COUNT)
        == expected_result[INDEXED_DOCUMENT_COUNT]
    )
    assert (
        sink.counters.get(INDEXED_DOCUMENT_VOLUME)
        == expected_result[INDEXED_DOCUMENT_VOLUME]
    )
    assert (
        sink.counters.get(DELETED_DOCUMENT_COUNT)
        == expected_result[DELETED_DOCUMENT_COUNT]
    )


@pytest.mark.asyncio
async def test_batch_bulk_with_retry():
    config = {
        "username": "elastic",
        "password": "changeme",
        "host": "http://nowhere.com:9200",
    }
    client = ESManagementClient(config)
    client.client = AsyncMock()
    sink = Sink(
        client=client,
        queue=None,
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )

    with mock.patch.object(asyncio, "sleep"):
        # first call raises exception, and the second call succeeds
        error_meta = Mock()
        error_meta.status = 429
        first_call_error = ApiError(429, meta=error_meta, body="error")
        second_call_result = {"items": []}
        client.client.bulk = AsyncMock(
            side_effect=[first_call_error, second_call_result]
        )
        await sink._batch_bulk([], {OP_INDEX: {}, OP_UPDATE: {}, OP_DELETE: {}})

        assert client.client.bulk.await_count == 2


@pytest.mark.asyncio
async def test_batch_bulk_with_error_monitor():
    config = {
        "username": "elastic",
        "password": "changeme",
        "host": "http://nowhere.com:9200",
    }
    client = ESManagementClient(config)
    client.client = AsyncMock()
    sink = Sink(
        client=client,
        queue=None,
        error_monitor=ErrorMonitor(max_consecutive_errors=10),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )

    def _create_response(num_successful, num_noop, num_failed):
        call_result = {"items": []}
        for id_ in range(num_successful):
            call_result["items"].append(
                {OP_UPDATE: {"_id": str(id_), "result": "updated"}}
            )

        for id_ in range(num_noop):
            call_result["items"].append(
                {OP_UPDATE: {"_id": str(id_ + num_successful), "result": "noop"}}
            )

        for id_ in range(num_failed):
            call_result["items"].append(
                {
                    OP_UPDATE: {
                        "_id": str(id_ + num_successful + num_noop),
                        "result": "failed",
                        "error": "something went wrong",
                    }
                }
            )

        return call_result

    # First add 40 successful items, 5 noop and 5 failed
    first_call_result = _create_response(40, 5, 5)

    # Then add 30 successful items, 10 noop and 10 failed
    second_call_result = _create_response(30, 10, 10)

    # Lastly add 0 successful items and 1 failed
    # This should fail the test as max consecutive_error_count=10
    # 10 errors came last in second call, and this one just triggers it
    third_call_result = _create_response(0, 0, 1)

    client.client.bulk = AsyncMock(
        side_effect=[first_call_result, second_call_result, third_call_result]
    )

    # call 1
    await sink._batch_bulk([], {OP_INDEX: {}, OP_UPDATE: {}, OP_DELETE: {}})
    # call 2
    await sink._batch_bulk([], {OP_INDEX: {}, OP_UPDATE: {}, OP_DELETE: {}})

    with pytest.raises(TooManyErrors):
        # call 3
        await sink._batch_bulk([], {OP_INDEX: {}, OP_UPDATE: {}, OP_DELETE: {}})


@pytest.mark.asyncio
async def test_batch_bulk_still_counts_successful_docs_when_error_monitor_raises():
    """When the error monitor raises mid-response, docs ES accepted in the same
    batch must still be reflected in `INDEXED_DOCUMENT_COUNT`."""
    config = {
        "username": "elastic",
        "password": "changeme",
        "host": "http://nowhere.com:9200",
    }
    client = ESManagementClient(config)
    client.client = AsyncMock()
    sink = Sink(
        client=client,
        queue=None,
        error_monitor=ErrorMonitor(max_consecutive_errors=10),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )

    # 11 failures (one over the threshold) followed by 5 ES-accepted successes.
    items = []
    for i in range(11):
        items.append({OP_INDEX: {"_id": f"fail_{i}", "error": "something went wrong"}})
    for i in range(5):
        items.append({OP_INDEX: {"_id": f"ok_{i}", "result": "created"}})

    client.bulk_insert = AsyncMock(return_value={"items": items, "errors": True})

    stats = {
        OP_INDEX: {f"fail_{i}": 100 for i in range(11)},
        OP_UPDATE: {},
        OP_DELETE: {},
    }
    for i in range(5):
        stats[OP_INDEX][f"ok_{i}"] = 50

    with pytest.raises(TooManyErrors):
        await sink._batch_bulk([], stats)

    assert sink.counters.get(INDEXED_DOCUMENT_COUNT) == 5
    assert sink.counters.get(INDEXED_DOCUMENT_VOLUME) == 5 * 50


def test_sync_orchestrator_passes_error_monitor_config_as_kwargs():
    """The `bulk.error_monitor` config block must reach individual `ErrorMonitor`
    fields; passing it positionally used to bind the whole dict to `enabled`
    and leave every threshold at its default."""
    config = {
        "host": "http://nowhere.com:9200",
        "username": "elastic",
        "password": "changeme",
        "bulk": {
            "error_monitor": {
                "enabled": True,
                "max_total_errors": 9999,
                "max_consecutive_errors": 42,
                "max_error_rate": 0.5,
                "error_window_size": 50,
                "error_queue_size": 7,
            }
        },
    }

    orchestrator = SyncOrchestrator(config)

    assert orchestrator.error_monitor.enabled is True
    assert orchestrator.error_monitor.max_total_errors == 9999
    assert orchestrator.error_monitor.max_consecutive_errors == 42
    assert orchestrator.error_monitor.max_error_rate == 0.5
    assert orchestrator.error_monitor.error_window_size == 50
    assert orchestrator.error_monitor.error_queue_size == 7


@pytest.mark.asyncio
async def test_batch_bulk_with_errors(patch_logger):
    config = {
        "username": "elastic",
        "password": "changeme",
        "host": "http://nowhere.com:9200",
    }
    client = ESManagementClient(config)
    client.client = AsyncMock()
    sink = Sink(
        client=client,
        queue=None,
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )

    with mock.patch.object(asyncio, "sleep"):
        error = {
            "type": "illegal_argument_exception",
            "reason": "if _id is specified it must not be empty",
        }
        mock_result = {
            "errors": True,
            "items": [{"index": {"_id": "1", "error": error}}],
        }
        client.client.bulk = AsyncMock(return_value=mock_result)
        await sink._batch_bulk([], {OP_INDEX: {"1": 20}, OP_UPDATE: {}, OP_DELETE: {}})
        patch_logger.assert_present(f"operation index failed for doc 1, {error}")


@patch("connectors.es.sink.CANCELATION_TIMEOUT", -1)
@pytest.mark.parametrize(
    "extractor_task, extractor_task_done, sink_task, sink_task_done, expected_result",
    [
        (None, False, None, False, True),
        (Mock(), False, None, False, False),
        (Mock(), True, None, False, True),
        (None, False, Mock(), False, False),
        (None, False, Mock(), True, True),
        (Mock(), False, Mock(), False, False),
        (Mock(), False, Mock(), True, False),
        (Mock(), True, Mock(), False, False),
        (Mock(), True, Mock(), True, True),
    ],
)
@pytest.mark.asyncio
async def test_sync_orchestrator_done_and_cleanup(
    extractor_task, extractor_task_done, sink_task, sink_task_done, expected_result
):
    if extractor_task is not None:
        extractor_task.cancel = Mock()
        extractor_task.done.return_value = extractor_task_done
    if sink_task is not None:
        sink_task.cancel = Mock()
        sink_task.done.return_value = sink_task_done

    config = {"host": "http://nowhere.com:9200", "user": "tarek", "password": "blah"}
    es = SyncOrchestrator(config)
    es._extractor = Mock()
    es._extractor.error = None
    es._sink = Mock()
    es._sink.error = None
    es._extractor_task = extractor_task
    es._sink_task = sink_task

    assert es.done() == expected_result

    await es.close()

    if extractor_task and not extractor_task_done:
        extractor_task.cancel.assert_called_once()
    if sink_task and not sink_task_done:
        sink_task.cancel.assert_called_once()


@pytest.mark.asyncio
async def test_extractor_put_doc():
    doc = {"id": 123}
    queue = Mock()
    queue.put = AsyncMock()
    extractor = Extractor(None, queue, INDEX)

    await extractor.put_doc(doc)
    queue.put.assert_awaited_once_with(doc)


@pytest.mark.asyncio
@mock.patch(
    "connectors.es.management_client.ESManagementClient.yield_existing_documents_metadata"
)
@mock.patch("connectors.utils.ConcurrentTasks.cancel")
async def test_extractor_get_docs_when_downloads_fail(
    yield_existing_documents_metadata, concurrent_tasks_cancel
):
    queue = await queue_mock()

    yield_existing_documents_metadata.return_value = AsyncIterator([])

    docs_from_source = [
        (DOC_ONE, crashing_lazy_download_fake(), "index"),
        (DOC_TWO, crashing_lazy_download_fake(), "index"),
        (DOC_THREE, crashing_lazy_download_fake(), "index"),
        (DOC_FOUR, crashing_lazy_download_fake(), "index"),
    ]
    # deep copying docs is needed as get_docs mutates the document ids which has side effects on other test
    # instances
    doc_generator = AsyncIterator([deepcopy(doc) for doc in docs_from_source])

    extractor = await setup_extractor(queue, content_extraction_enabled=True)

    await extractor.run(doc_generator, JobType.FULL)
    concurrent_tasks_cancel.assert_called_once()


@pytest.mark.asyncio
async def test_force_canceled_extractor_put_doc():
    doc = {"id": 123}
    queue = Mock()
    queue.put = AsyncMock()
    extractor = Extractor(None, queue, INDEX)

    extractor.force_cancel()
    with pytest.raises(ForceCanceledError):
        await extractor.put_doc(doc)
        queue.put.assert_not_awaited()


@pytest.mark.asyncio
async def test_force_canceled_extractor_with_other_errors(patch_logger):
    queue = Mock()
    queue.put = AsyncMock()
    extractor = Extractor(
        None,
        queue,
        INDEX,
    )
    generator = AsyncMock(side_effect=Exception("a non-ForceCanceledError"))

    extractor.force_cancel()
    # no error thrown
    await extractor.run(generator, JobType.FULL)
    patch_logger.assert_present(
        "Extractor did not stop within 5 seconds of cancelation, force-canceling the task."
    )


@pytest.mark.asyncio
async def test_sink_fetch_doc():
    expected_doc = {"id": 123}
    queue = Mock()
    queue.get = AsyncMock(return_value=expected_doc)
    sink = Sink(
        None,
        queue,
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )

    doc = await sink.fetch_doc()
    queue.get.assert_awaited_once()
    assert doc == expected_doc


@pytest.mark.asyncio
async def test_force_canceled_sink_fetch_doc():
    expected_doc = {"id": 123}
    queue = Mock()
    queue.get = AsyncMock(return_value=expected_doc)
    sink = Sink(
        None,
        queue,
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )

    sink.force_cancel()
    with pytest.raises(ForceCanceledError):
        await sink.fetch_doc()
        queue.get.assert_not_awaited()


@pytest.mark.asyncio
async def test_force_canceled_sink_with_other_errors(patch_logger):
    queue = Mock()
    queue.get = AsyncMock(side_effect=Exception("a non-ForceCanceledError"))
    sink = Sink(
        None,
        queue,
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
    )

    sink.force_cancel()
    # no error thrown
    await sink.run()
    patch_logger.assert_present(
        "Sink did not stop within 5 seconds of cancelation, force-canceling the task."
    )


@pytest.mark.parametrize(
    "extractor_task_done, sink_task_done, force_cancel",
    [
        (
            itertools.chain([False, False, False], itertools.repeat(True)),
            itertools.chain([False, False], itertools.repeat(True)),
            False,
        ),
        (
            itertools.chain([False, False, False], itertools.repeat(True)),
            itertools.repeat(False),
            True,
        ),
        (
            itertools.repeat(False),
            itertools.chain([False, False], itertools.repeat(True)),
            True,
        ),
        (
            itertools.repeat(False),
            itertools.repeat(False),
            True,
        ),
    ],
)
@pytest.mark.asyncio
async def test_cancel_sync(extractor_task_done, sink_task_done, force_cancel):
    config = {"host": "http://nowhere.com:9200", "user": "tarek", "password": "blah"}
    es = SyncOrchestrator(config)
    es._extractor = Mock()
    es._extractor.force_cancel = Mock()

    es._sink = Mock()
    es._sink.force_cancel = Mock()

    es._extractor_task = mock.create_autospec(asyncio.Task)
    es._extractor_task.cancel = Mock()
    es._extractor_task.done = Mock(side_effect=extractor_task_done)

    es._sink_task = mock.create_autospec(asyncio.Task)
    es._sink_task.cancel = Mock()
    es._sink_task.done = Mock(side_effect=sink_task_done)

    with mock.patch.object(asyncio, "sleep"):
        await es.cancel()
        es._extractor_task.cancel.assert_called_once()
        es._sink_task.cancel.assert_called_once()

        if force_cancel:
            es._extractor.force_cancel.assert_called_once()
            es._sink.force_cancel.assert_called_once()
        else:
            es._extractor.force_cancel.assert_not_called()
            es._sink.force_cancel.assert_not_called()


@pytest.mark.asyncio
async def test_extractor_run_when_mem_full_is_raised():
    docs_from_source = [
        {"_id": 1},
        {"_id": 2},
        {"_id": 3},
    ]
    es_client = AsyncMock()
    es_client.yield_existing_documents_metadata = Mock(return_value=AsyncIterator([]))

    queue = await queue_mock()
    queue.clear = Mock()

    def _put_side_effect(value):
        if isinstance(value, str):
            pass
        else:
            raise asyncio.QueueFull()

    queue.put = AsyncMock(side_effect=_put_side_effect)
    # deep copying docs is needed as get_access_control_docs mutates the document ids which has side effects on other test
    # instances
    doc_generator = AsyncIterator([(doc, None, "") for doc in docs_from_source])

    extractor = Extractor(es_client, queue, INDEX)

    await extractor.run(doc_generator, JobType.FULL)

    queue.clear.assert_called_once()
    assert isinstance(extractor.error, ElasticsearchOverloadedError)


@pytest.mark.asyncio
async def test_should_not_log_bulk_operations_if_doc_id_tracing_is_disabled(
    patch_logger,
):
    action = "create"
    result = "created"
    operations = []
    client = Mock()
    client.bulk_insert = AsyncMock(
        return_value={"items": [successful_bulk_action(DOC_ONE_ID, action, result)]}
    )
    sink = Sink(
        client=client,
        queue=Mock(),
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
        enable_bulk_operations_logging=False,
    )

    await sink._batch_bulk(operations, STATS)

    patch_logger.assert_not_present(
        successful_action_log_message(DOC_ONE_ID, action, result)
    )


@pytest.mark.parametrize(
    "operation_results, expected_logs",
    [
        (
            [successful_bulk_action(DOC_ONE_ID, "create", "created")],
            [successful_action_log_message(DOC_ONE_ID, "create", "created")],
        ),
        (
            [successful_bulk_action(DOC_ONE_ID, "delete", "deleted")],
            [successful_action_log_message(DOC_ONE_ID, "delete", "deleted")],
        ),
        (
            [successful_bulk_action(DOC_ONE_ID, "index", "created")],
            [successful_action_log_message(DOC_ONE_ID, "index", "created")],
        ),
        (
            [successful_bulk_action(DOC_ONE_ID, "update", "updated")],
            [successful_action_log_message(DOC_ONE_ID, "update", "updated")],
        ),
        (
            [successful_bulk_action(DOC_ONE_ID, "delete", "not_found")],
            [
                successful_operation_with_non_successful_result_log_message(
                    DOC_ONE_ID, "delete", "not_found"
                )
            ],
        ),
        (
            [failed_bulk_action(DOC_ONE_ID, "create", "not_found")],
            [failed_action_log_message(DOC_ONE_ID, "create", "not_found")],
        ),
        (
            [failed_bulk_action(DOC_ONE_ID, "delete", "not_found")],
            [failed_action_log_message(DOC_ONE_ID, "delete", "not_found")],
        ),
        (
            [failed_bulk_action(DOC_ONE_ID, "index", "not_found")],
            [failed_action_log_message(DOC_ONE_ID, "index", "not_found")],
        ),
        (
            [failed_bulk_action(DOC_ONE_ID, "update", "not_found")],
            [failed_action_log_message(DOC_ONE_ID, "update", "not_found")],
        ),
        (
            [
                successful_bulk_action(DOC_ONE_ID, "create", "created"),
                successful_bulk_action(DOC_TWO_ID, "update", "updated"),
                failed_bulk_action(DOC_THREE_ID, "update", "not_found"),
            ],
            [
                successful_action_log_message(DOC_ONE_ID, "create", "created"),
                successful_action_log_message(DOC_TWO_ID, "update", "updated"),
                failed_action_log_message(DOC_THREE_ID, "update", "not_found"),
            ],
        ),
    ],
)
@pytest.mark.asyncio
async def test_should_log_bulk_operations_if_doc_id_tracing_is_enabled(
    patch_logger, operation_results, expected_logs
):
    operations = []
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": operation_results})
    sink = Sink(
        client=client,
        queue=Mock(),
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
        enable_bulk_operations_logging=True,
    )

    await sink._batch_bulk(operations, STATS)

    for log in expected_logs:
        patch_logger.assert_present(log)


@pytest.mark.asyncio
async def test_should_log_error_when_id_is_missing(patch_logger):
    operations = []
    client = Mock()
    # item missing id
    item = {"index": {"result": "created"}}

    client.bulk_insert = AsyncMock(
        return_value={
            "items": [item, successful_bulk_action(DOC_ONE_ID, "create", "created")]
        }
    )
    sink = Sink(
        client=client,
        queue=Mock(),
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
        enable_bulk_operations_logging=True,
    )

    await sink._batch_bulk(operations, STATS)

    patch_logger.assert_present(f"Could not retrieve '_id' for document {item}")

    # Should continue logging after a failure
    patch_logger.assert_present(
        successful_action_log_message(DOC_ONE_ID, "create", "created")
    )


@pytest.mark.asyncio
async def test_should_log_error_when_unknown_action_item_returned(patch_logger):
    operations = []
    client = Mock()

    # item with unknown action item
    item = {"unknown": {"result": "created"}}

    client.bulk_insert = AsyncMock(
        return_value={
            "items": [item, successful_bulk_action(DOC_ONE_ID, "create", "created")]
        }
    )
    sink = Sink(
        client=client,
        queue=Mock(),
        error_monitor=Mock(),
        chunk_size=0,
        pipeline={"name": "pipeline"},
        chunk_mem_size=0,
        max_concurrency=0,
        max_retries=3,
        retry_interval=10,
        enable_bulk_operations_logging=True,
    )

    await sink._batch_bulk(operations, STATS)

    patch_logger.assert_present(
        f"Unknown action item returned from _bulk API for item {item}"
    )

    # Should continue logging after a failure
    patch_logger.assert_present(
        successful_action_log_message(DOC_ONE_ID, "create", "created")
    )


MIB = 1024 * 1024


def _make_index_doc(doc_id):
    return {
        "_op_type": OP_INDEX,
        "_index": INDEX,
        "_id": str(doc_id),
        "doc": {"id": str(doc_id), "_timestamp": TIMESTAMP.isoformat()},
    }


def _make_delete_doc(doc_id):
    return {
        "_op_type": OP_DELETE,
        "_index": INDEX,
        "_id": str(doc_id),
    }


def _doc(doc_id, op=OP_INDEX):
    """Build the minimal doc shape `Sink._run` reads: `_op_type`/`_index`/`_id`
    metadata at the top level, plus a `doc` body for index/update."""
    doc = {"_op_type": op, "_index": INDEX, "_id": doc_id}
    if op != OP_DELETE:
        doc["doc"] = {"id": doc_id}
    return doc


def _doc_ids(batch):
    """Return the doc IDs in a dispatched bulk batch, in original order.

    A bulk batch is a flat list of action/source lines. Each doc contributes
    exactly one action line (`{"index"|"update"|"delete": {"_id": ...}}`),
    so we pull `_id` off those and skip the body payload lines."""
    return [
        entry[op]["_id"]
        for entry in batch
        for op in (OP_INDEX, OP_UPDATE, OP_DELETE)
        if op in entry
    ]


def _queue_yielding(items):
    """Mock queue whose `get` yields the given `(size, doc)` items in order,
    then `END_DOCS` so `Sink._run` exits cleanly."""
    queue = Mock()
    queue.get = AsyncMock(side_effect=[*items, (0, END_DOCS)])
    return queue


def _make_cap_sink(client, queue, max_text_document_size=None):
    return Sink(
        client=client,
        queue=queue,
        error_monitor=Mock(),
        chunk_size=1000,
        pipeline={"name": "pipeline"},
        chunk_mem_size=100,
        max_concurrency=1,
        max_retries=3,
        retry_interval=10,
        max_text_document_size=max_text_document_size,
    )


def _inflate_body(doc, payload_bytes):
    """Pad doc['doc']['body'] so that the bulk-op pair serializes above payload_bytes."""
    doc["doc"]["body"] = "x" * payload_bytes
    return doc


@pytest.mark.asyncio
async def test_sink_drops_doc_exceeding_max_text_document_size(patch_logger):
    big_doc = _inflate_body(_make_index_doc(DOC_ONE_ID), 5 * 1024 * 1024)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(5 * 1024 * 1024, big_doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    # Match `Sink._run`'s wire-byte measurement: reuse the same
    # elastic_transport `JsonSerializer` it uses in production.
    serializer = JsonSerializer()
    expected_serialized = sum(
        len(serializer.json_dumps(op)) for op in sink._bulk_op(big_doc, OP_INDEX)
    )
    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 1
    client.bulk_insert.assert_not_awaited()
    patch_logger.assert_present(
        f"Dropping doc id={big_doc['_id']} index={INDEX} op={OP_INDEX}: "
        f"serialized text size {expected_serialized}B exceeds "
        f"elasticsearch.bulk.max_text_document_size "
        f"({3 * 1024 * 1024}B)"
    )


@pytest.mark.asyncio
async def test_sink_does_not_drop_doc_within_max_text_document_size():
    small_doc = _make_index_doc(DOC_ONE_ID)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(1024, small_doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 0
    client.bulk_insert.assert_awaited_once()


@pytest.mark.parametrize("bad_value", [-1, -3, -1024])
def test_sink_rejects_negative_max_text_document_size(bad_value):
    # Without explicit validation, a negative cap would be truthy and every
    # serialized_size > cap, so every non-attachment doc would be dropped
    # silently. Refuse the config at construction instead.
    queue = _queue_yielding([])
    with pytest.raises(ValueError, match="max_text_document_size must be >= 0"):
        _make_cap_sink(Mock(), queue, max_text_document_size=bad_value)


@pytest.mark.asyncio
@pytest.mark.parametrize("disabled_value", [None, 0])
async def test_sink_max_text_document_size_disabled(disabled_value):
    big_doc = _inflate_body(_make_index_doc(DOC_ONE_ID), 5 * 1024 * 1024)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(10 * 1024 * 1024, big_doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=disabled_value)

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 0
    client.bulk_insert.assert_awaited_once()


@pytest.mark.asyncio
async def test_sink_does_not_drop_delete_op_even_if_oversized():
    # DELETE ops carry no body; they should never be dropped by the cap.
    delete_doc = _make_delete_doc(DOC_ONE_ID)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(10 * 1024 * 1024, delete_doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 0
    client.bulk_insert.assert_awaited_once()


@pytest.mark.asyncio
async def test_sink_drops_only_oversized_doc_in_mixed_batch():
    doc_small_1 = _make_index_doc(DOC_ONE_ID)
    doc_big = _inflate_body(_make_index_doc(DOC_TWO_ID), 5 * 1024 * 1024)
    doc_small_2 = _make_index_doc(DOC_THREE_ID)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding(
        [
            (1024, doc_small_1),
            (10 * 1024 * 1024, doc_big),
            (1024, doc_small_2),
        ]
    )

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 1
    client.bulk_insert.assert_awaited_once()
    sent_ops = client.bulk_insert.await_args.args[0]
    sent_ids = {
        op[OP_INDEX]["_id"]
        for op in sent_ops
        if isinstance(op, dict) and OP_INDEX in op
    }
    assert sent_ids == {str(DOC_ONE_ID), str(DOC_THREE_ID)}


@pytest.mark.asyncio
async def test_sink_does_not_drop_doc_with_attachment_even_if_oversized():
    # Binary attachments are governed by service.max_file_download_size,
    # not by elasticsearch.bulk.max_text_document_size. Even if the doc's
    # serialized size exceeds the cap, presence of `_attachment` exempts it.
    big_attachment_doc = _make_index_doc(DOC_ONE_ID)
    big_attachment_doc["doc"]["_attachment"] = "x" * (5 * 1024 * 1024)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(10 * 1024 * 1024, big_attachment_doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 0
    client.bulk_insert.assert_awaited_once()


@pytest.mark.asyncio
@pytest.mark.parametrize("attachment_value", ["", None, 0, False, []])
async def test_sink_attachment_gate_uses_key_presence_not_value(attachment_value):
    # The cap exempts any doc carrying an `_attachment` key, regardless of
    # the value. Connectors only ever set `_attachment` to a base64 string
    # for real binary content, but the gate must not drop docs that pass an
    # empty/None/falsy `_attachment` -- those still belong on the binary path.
    doc = _inflate_body(_make_index_doc(DOC_ONE_ID), 5 * 1024 * 1024)
    doc["doc"]["_attachment"] = attachment_value
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(10 * 1024 * 1024, doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 0
    client.bulk_insert.assert_awaited_once()


@pytest.mark.asyncio
async def test_sink_drops_oversized_doc_when_attachment_key_absent():
    # Mirror of the gate from the other side: oversized doc with NO
    # `_attachment` key at all (neither set to None, nor "") goes through
    # the text-cap path and is dropped.
    doc = _inflate_body(_make_index_doc(DOC_ONE_ID), 5 * 1024 * 1024)
    assert "_attachment" not in doc["doc"]  # sanity check
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(10 * 1024 * 1024, doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 1
    client.bulk_insert.assert_not_awaited()


@pytest.mark.asyncio
async def test_sink_drops_doc_with_body_only_when_oversized():
    # DES-text shape: doc has `body` but no `_attachment` -> cap applies.
    doc = _inflate_body(_make_index_doc(DOC_ONE_ID), 5 * 1024 * 1024)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(5 * 1024 * 1024, doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 1
    client.bulk_insert.assert_not_awaited()


@pytest.mark.asyncio
async def test_sink_drops_structured_only_doc_when_oversized():
    # Structured-only shape: no `_attachment`, no `body`. Cap still applies.
    doc = _make_index_doc(DOC_ONE_ID)
    doc["doc"]["title"] = "x" * (5 * 1024 * 1024)
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(5 * 1024 * 1024, doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 1
    client.bulk_insert.assert_not_awaited()


@pytest.mark.asyncio
async def test_sink_measures_serialized_size_in_wire_utf8_bytes():
    # Regression: the cap must be measured against the bytes the ES client
    # actually sends, which is `ensure_ascii=False` + UTF-8. Building a body
    # under the cap in wire bytes but over it in default `json.dumps` chars
    # (e.g. each `é` is 2 wire bytes vs 6 ASCII-escaped chars) must NOT drop.
    #
    # cap = 1 MiB. Pick a count of `é` chars whose ASCII-escape (`\u00e9`,
    # 6 chars) exceeds the cap, but whose UTF-8 encoding (2 bytes) does not.
    # 200_000 * 6 = 1_200_000 bytes (over) vs 200_000 * 2 = 400_000 bytes
    # (well under), so the wire-byte cap must let this through.
    cap_mib = 1
    doc = _make_index_doc(DOC_ONE_ID)
    doc["doc"]["body"] = "é" * 200_000
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(1024, doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=cap_mib)

    await sink.run()

    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 0
    client.bulk_insert.assert_awaited_once()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "field_value",
    [
        datetime.datetime(2024, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc),
        datetime.date(2024, 1, 2),
        uuid.UUID("12345678-1234-5678-1234-567812345678"),
        Decimal("3.14"),
    ],
    ids=["datetime", "date", "uuid", "decimal"],
)
async def test_sink_cap_handles_non_json_native_types(field_value):
    # Regression: connectors (e.g. the directory source) may yield docs
    # containing types that bare `json.dumps` can't serialize. The cap path
    # must use the same `default()` handling as the wire serializer so it
    # doesn't crash the sink before any doc is indexed.
    doc = _make_index_doc(DOC_ONE_ID)
    doc["doc"]["timestamp_field"] = field_value
    client = Mock()
    client.bulk_insert = AsyncMock(return_value={"items": []})
    queue = _queue_yielding([(1024, doc)])

    sink = _make_cap_sink(client, queue, max_text_document_size=3)  # MiB

    await sink.run()

    assert sink.error is None
    assert sink.counters.get(DOCS_DROPPED_TOO_LARGE) == 0
    client.bulk_insert.assert_awaited_once()


def _make_sink(queue, *, chunk_size, chunk_mem_size):
    sink = Sink(
        client=None,
        queue=queue,
        chunk_size=chunk_size,
        pipeline={"name": "pipeline"},
        chunk_mem_size=chunk_mem_size,
        max_concurrency=2,
        max_retries=3,
        retry_interval=0,
        error_monitor=Mock(),
    )
    sink._batch_bulk = AsyncMock(return_value={"items": []})
    return sink


def _dispatched_batches(sink):
    """All batches the sink ever dispatched (in-loop + trailing flush)."""
    return [call_args.args[0] for call_args in sink._batch_bulk.await_args_list]


def _capture_in_loop_dispatch_names(sink):
    """Spy on `sink.bulk_tasks.put` so tests can tell in-loop dispatches
    (which go through `put`) apart from the trailing END_DOCS flush (which
    calls `_batch_bulk` directly), and assert on the bulk-request names."""
    names = []

    async def _put(callable_, name):
        names.append(name)
        await callable_()

    sink.bulk_tasks.put = _put
    return names


@pytest.mark.asyncio
async def test_sink_flushes_before_overflowing_chunk_mem_size():
    # Original bug: 5 docs at 1/1/1/1/2 MiB with chunk_mem_size=5 used to
    # dispatch a single 6 MiB batch. The pre-flush must ship the first 4
    # docs together (4 MiB) and the trailing flush ships the 2 MiB doc alone.
    items = [
        (1 * MIB, _doc("1")),
        (1 * MIB, _doc("2")),
        (1 * MIB, _doc("3")),
        (1 * MIB, _doc("4")),
        (2 * MIB, _doc("5")),
    ]
    sink = _make_sink(_queue_yielding(items), chunk_size=1000, chunk_mem_size=5)

    await sink._run()

    assert [_doc_ids(b) for b in _dispatched_batches(sink)] == [
        ["1", "2", "3", "4"],
        ["5"],
    ]


@pytest.mark.asyncio
async def test_sink_flushes_before_overflowing_chunk_size_with_mixed_ops():
    # `_bulk_op` emits 1 entry per delete and 2 per index/update. With a
    # naive `len(batch) >= chunk_size` check, `delete + index + update` would
    # grow the batch to 5 entries before any flush -- exceeding chunk_size=4.
    # The prospective pre-flush must split the run so every batch stays <= 4.
    items = [
        (1, _doc("1", OP_DELETE)),
        (1, _doc("2")),
        (1, _doc("3", OP_UPDATE)),
        (1, _doc("4")),
    ]
    sink = _make_sink(_queue_yielding(items), chunk_size=4, chunk_mem_size=1024)

    await sink._run()

    # First batch: delete + index = 1 + 2 = 3 entries.
    # Second batch: update + index = 2 + 2 = 4 entries (exactly chunk_size).
    assert [_doc_ids(b) for b in _dispatched_batches(sink)] == [
        ["1", "2"],
        ["3", "4"],
    ]


@pytest.mark.asyncio
async def test_sink_dispatches_full_batch_without_waiting_for_next_doc():
    # When a doc fills the batch exactly to chunk_size, it must be shipped
    # immediately via `bulk_tasks.put` -- not held until the next doc or the
    # trailing END_DOCS flush. Also pins that the dispatched-task names
    # count batches (#1, #2), not fetched docs.
    items = [(1, _doc(str(i))) for i in range(1, 5)]  # 4 docs -> 2 full batches of 4
    sink = _make_sink(_queue_yielding(items), chunk_size=4, chunk_mem_size=1024)
    in_loop_names = _capture_in_loop_dispatch_names(sink)

    await sink._run()

    assert in_loop_names == [
        "Elasticsearch Sink: _bulk batch #1",
        "Elasticsearch Sink: _bulk batch #2",
    ]
    assert [_doc_ids(b) for b in _dispatched_batches(sink)] == [
        ["1", "2"],
        ["3", "4"],
    ]


@pytest.mark.asyncio
async def test_sink_dispatches_oversized_single_doc_alone():
    # A single 10 MiB doc against a 5 MiB cap has no batch to split it from.
    # The pre-flush must skip (the `if batch` guard prevents an empty bulk
    # dispatch) and the post-append flush must still ship the doc, so the
    # loop doesn't stall on a bigger-than-cap document.
    sink = _make_sink(
        _queue_yielding([(10 * MIB, _doc("big"))]),
        chunk_size=1000,
        chunk_mem_size=5,
    )
    in_loop_names = _capture_in_loop_dispatch_names(sink)

    await sink._run()

    assert in_loop_names == ["Elasticsearch Sink: _bulk batch #1"]
    assert [_doc_ids(b) for b in _dispatched_batches(sink)] == [["big"]]
