import json
from typing import Any, Dict, List, Optional, cast
from unittest.mock import MagicMock, patch

import pytest

from acryl_datahub_cloud.datahub_context_doc_proposals.context_doc_proposals_source import (
    ADMIN_ROLE_URN,
    DOMAINS_SP_FIELD,
    EXPERTS_SP_FIELD,
    MAX_DOCS_PER_PROPOSAL,
    PUBLISHED_STAGE_URN,
    SYSTEM_URN,
    DataHubContextDocProposalsSource,
    DataHubContextDocProposalsSourceConfig,
    DocInfo,
    _collect_expert_urns,
    _group_docs_into_batches,
    _grouping_key,
    _parse_json_list_field,
    _parse_structured_properties,
    _split_user_and_group_urns,
)
from acryl_datahub_cloud.metadata.schema_classes import (
    ActionRequestInfoClass,
    ActionRequestParamsClass,
    DocumentChangeClass,
    DocumentProposalTypeClass,
    DocumentsProposalClass,
)
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.graph.client import DataHubGraph

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def mock_graph() -> MagicMock:
    return MagicMock(spec=DataHubGraph)


@pytest.fixture
def mock_context(mock_graph: MagicMock) -> MagicMock:
    ctx = MagicMock(spec=PipelineContext)
    ctx.require_graph.return_value = mock_graph
    return ctx


@pytest.fixture
def source(mock_context: MagicMock) -> DataHubContextDocProposalsSource:
    src = DataHubContextDocProposalsSource(
        DataHubContextDocProposalsSourceConfig(), mock_context
    )
    src.graph = cast(DataHubGraph, mock_context.require_graph())
    return src


# ---------------------------------------------------------------------------
# JSON helpers — mirrors the wire format written by the enrichment pipeline
# ---------------------------------------------------------------------------


def _expert_json(principal_urn: str, score: float = 0.5) -> str:
    return json.dumps(
        {
            "principal_urn": principal_urn,
            "principal_type": "corpuser",
            "score": score,
            "dataset_count": 3,
            "domain_count": 1,
            "reasoning": "test",
        }
    )


def _domain_json(domain_urn: str, score: float = 0.8) -> str:
    return json.dumps(
        {
            "domain_urn": domain_urn,
            "score": score,
            "dataset_count": 5,
            "reasoning": "test",
        }
    )


# ---------------------------------------------------------------------------
# _parse_json_list_field
# ---------------------------------------------------------------------------


class TestParseJsonListField:
    def test_array_of_json_strings(self) -> None:
        raw = json.dumps([_expert_json("urn:li:corpuser:alice")])
        result = _parse_json_list_field(raw, "principal_urn")
        assert result == ["urn:li:corpuser:alice"]

    def test_array_of_dicts(self) -> None:
        raw = json.dumps([{"principal_urn": "urn:li:corpuser:bob", "score": 0.9}])
        result = _parse_json_list_field(raw, "principal_urn")
        assert result == ["urn:li:corpuser:bob"]

    def test_multiple_entries(self) -> None:
        entries = [
            _expert_json("urn:li:corpuser:alice"),
            _expert_json("urn:li:corpuser:bob"),
        ]
        raw = json.dumps(entries)
        result = _parse_json_list_field(raw, "principal_urn")
        assert result == ["urn:li:corpuser:alice", "urn:li:corpuser:bob"]

    def test_single_object(self) -> None:
        raw = _expert_json("urn:li:corpuser:carol")
        result = _parse_json_list_field(raw, "principal_urn")
        assert result == ["urn:li:corpuser:carol"]

    def test_invalid_json(self) -> None:
        assert _parse_json_list_field("not-json{", "principal_urn") == []

    def test_empty_string(self) -> None:
        assert _parse_json_list_field("", "principal_urn") == []

    def test_missing_key(self) -> None:
        raw = json.dumps([{"other_field": "val"}])
        assert _parse_json_list_field(raw, "principal_urn") == []


# ---------------------------------------------------------------------------
# _parse_structured_properties
# ---------------------------------------------------------------------------


class TestParseStructuredProperties:
    def test_both_experts_and_domains(self) -> None:
        extra_props: List[Dict[str, str]] = [
            {
                "name": EXPERTS_SP_FIELD,
                "value": json.dumps([_expert_json("urn:li:corpuser:alice")]),
            },
            {
                "name": DOMAINS_SP_FIELD,
                "value": json.dumps([_domain_json("urn:li:domain:finance")]),
            },
        ]
        experts, domains = _parse_structured_properties(extra_props)
        assert experts == ["urn:li:corpuser:alice"]
        assert domains == ["urn:li:domain:finance"]

    def test_experts_only(self) -> None:
        extra_props: List[Dict[str, str]] = [
            {
                "name": EXPERTS_SP_FIELD,
                "value": json.dumps([_expert_json("urn:li:corpuser:alice")]),
            },
        ]
        experts, domains = _parse_structured_properties(extra_props)
        assert experts == ["urn:li:corpuser:alice"]
        assert domains == []

    def test_no_relevant_props(self) -> None:
        extra_props: List[Dict[str, str]] = [
            {"name": "some.other.field", "value": '"hello"'}
        ]
        experts, domains = _parse_structured_properties(extra_props)
        assert experts == []
        assert domains == []

    def test_empty_value_skipped(self) -> None:
        extra_props: List[Dict[str, str]] = [{"name": EXPERTS_SP_FIELD, "value": ""}]
        experts, domains = _parse_structured_properties(extra_props)
        assert experts == []
        assert domains == []


# ---------------------------------------------------------------------------
# _grouping_key
# ---------------------------------------------------------------------------


class TestGroupingKey:
    def test_same_experts_same_key_regardless_of_domain(self) -> None:
        doc_a = DocInfo(
            urn="urn:li:document:a",
            expert_urns=["urn:li:corpuser:bob", "urn:li:corpuser:alice"],
            domain_urns=["urn:li:domain:fin"],
        )
        doc_b = DocInfo(
            urn="urn:li:document:b",
            expert_urns=["urn:li:corpuser:alice", "urn:li:corpuser:bob"],
            domain_urns=["urn:li:domain:eng"],
        )
        assert _grouping_key(doc_a) == _grouping_key(doc_b)

    def test_different_experts_different_key(self) -> None:
        doc_a = DocInfo(urn="urn:li:document:a", expert_urns=["urn:li:corpuser:alice"])
        doc_b = DocInfo(urn="urn:li:document:b", expert_urns=["urn:li:corpuser:bob"])
        assert _grouping_key(doc_a) != _grouping_key(doc_b)

    def test_no_experts_groups_by_domain(self) -> None:
        doc_a = DocInfo(urn="urn:li:document:a", domain_urns=["urn:li:domain:fin"])
        doc_b = DocInfo(urn="urn:li:document:b", domain_urns=["urn:li:domain:fin"])
        doc_c = DocInfo(urn="urn:li:document:c", domain_urns=["urn:li:domain:eng"])
        assert _grouping_key(doc_a) == _grouping_key(doc_b)
        assert _grouping_key(doc_a) != _grouping_key(doc_c)

    def test_no_experts_no_domains(self) -> None:
        doc_a = DocInfo(urn="urn:li:document:a")
        doc_b = DocInfo(urn="urn:li:document:b")
        assert _grouping_key(doc_a) == _grouping_key(doc_b)

    def test_experts_key_differs_from_domain_key(self) -> None:
        doc_with_expert = DocInfo(
            urn="urn:li:document:a", expert_urns=["urn:li:corpuser:alice"]
        )
        doc_with_domain = DocInfo(
            urn="urn:li:document:b", domain_urns=["urn:li:domain:fin"]
        )
        assert _grouping_key(doc_with_expert) != _grouping_key(doc_with_domain)


# ---------------------------------------------------------------------------
# _group_docs_into_batches
# ---------------------------------------------------------------------------


class TestGroupDocsIntoBatches:
    def test_groups_by_matching_key(self) -> None:
        docs = [
            DocInfo(urn=f"urn:li:document:{i}", expert_urns=["urn:li:corpuser:alice"])
            for i in range(3)
        ]
        batches = _group_docs_into_batches(docs)
        assert len(batches) == 1
        assert len(batches[0]) == 3

    def test_splits_large_groups(self) -> None:
        docs = [
            DocInfo(urn=f"urn:li:document:{i}", expert_urns=["urn:li:corpuser:alice"])
            for i in range(8)
        ]
        batches = _group_docs_into_batches(docs)
        assert len(batches) == 2
        assert len(batches[0]) == MAX_DOCS_PER_PROPOSAL
        assert len(batches[1]) == 3

    def test_different_keys_separate_batches(self) -> None:
        docs = [
            DocInfo(urn="urn:li:document:1", expert_urns=["urn:li:corpuser:alice"]),
            DocInfo(urn="urn:li:document:2", expert_urns=["urn:li:corpuser:bob"]),
            DocInfo(urn="urn:li:document:3", expert_urns=["urn:li:corpuser:alice"]),
        ]
        batches = _group_docs_into_batches(docs)
        assert len(batches) == 2
        alice_batch = [
            b for b in batches if b[0].expert_urns == ["urn:li:corpuser:alice"]
        ]
        bob_batch = [b for b in batches if b[0].expert_urns == ["urn:li:corpuser:bob"]]
        assert len(alice_batch) == 1
        assert len(alice_batch[0]) == 2
        assert len(bob_batch) == 1
        assert len(bob_batch[0]) == 1

    def test_empty_input(self) -> None:
        assert _group_docs_into_batches([]) == []

    def test_no_experts_grouped_together(self) -> None:
        docs = [DocInfo(urn=f"urn:li:document:{i}") for i in range(3)]
        batches = _group_docs_into_batches(docs)
        assert len(batches) == 1
        assert len(batches[0]) == 3


# ---------------------------------------------------------------------------
# _collect_expert_urns
# ---------------------------------------------------------------------------


class TestCollectExpertUrns:
    def test_deduplicates_across_docs(self) -> None:
        batch = [
            DocInfo(
                urn="urn:li:document:1",
                expert_urns=["urn:li:corpuser:alice", "urn:li:corpuser:bob"],
            ),
            DocInfo(
                urn="urn:li:document:2",
                expert_urns=["urn:li:corpuser:bob", "urn:li:corpuser:carol"],
            ),
        ]
        result = _collect_expert_urns(batch)
        assert result == [
            "urn:li:corpuser:alice",
            "urn:li:corpuser:bob",
            "urn:li:corpuser:carol",
        ]

    def test_empty_batch(self) -> None:
        assert _collect_expert_urns([]) == []

    def test_no_experts(self) -> None:
        assert _collect_expert_urns([DocInfo(urn="urn:li:document:1")]) == []


# ---------------------------------------------------------------------------
# _split_user_and_group_urns
# ---------------------------------------------------------------------------


class TestSplitUserAndGroupUrns:
    def test_users_only(self) -> None:
        users, groups = _split_user_and_group_urns(
            ["urn:li:corpuser:alice", "urn:li:corpuser:bob"]
        )
        assert users == ["urn:li:corpuser:alice", "urn:li:corpuser:bob"]
        assert groups == []

    def test_groups_only(self) -> None:
        users, groups = _split_user_and_group_urns(
            ["urn:li:corpGroup:eng", "urn:li:corpGroup:data"]
        )
        assert users == []
        assert groups == ["urn:li:corpGroup:eng", "urn:li:corpGroup:data"]

    def test_mixed(self) -> None:
        users, groups = _split_user_and_group_urns(
            [
                "urn:li:corpuser:alice",
                "urn:li:corpGroup:eng",
                "urn:li:corpuser:bob",
            ]
        )
        assert users == ["urn:li:corpuser:alice", "urn:li:corpuser:bob"]
        assert groups == ["urn:li:corpGroup:eng"]

    def test_empty(self) -> None:
        users, groups = _split_user_and_group_urns([])
        assert users == []
        assert groups == []


# ---------------------------------------------------------------------------
# _scroll_unpublished_docs
# ---------------------------------------------------------------------------


def _make_scroll_response(
    doc_entries: List[Dict[str, Any]],
    next_scroll_id: Optional[str] = None,
) -> Dict[str, Any]:
    search_results = []
    for entry in doc_entries:
        extra_props = []
        if "experts" in entry:
            extra_props.append(
                {
                    "name": EXPERTS_SP_FIELD,
                    "value": json.dumps(entry["experts"]),
                }
            )
        if "domains" in entry:
            extra_props.append(
                {
                    "name": DOMAINS_SP_FIELD,
                    "value": json.dumps(entry["domains"]),
                }
            )
        search_results.append(
            {
                "entity": {"urn": entry["urn"]},
                "extraProperties": extra_props,
            }
        )
    return {
        "scrollAcrossEntities": {
            "nextScrollId": next_scroll_id,
            "searchResults": search_results,
        }
    }


def test_scroll_unpublished_docs_single_page(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.return_value = _make_scroll_response(
        [
            {"urn": "urn:li:document:1"},
            {"urn": "urn:li:document:2"},
        ]
    )

    docs = source._scroll_unpublished_docs()

    assert [d.urn for d in docs] == ["urn:li:document:1", "urn:li:document:2"]
    assert source.report.unpublished_docs_found == 2


def test_scroll_unpublished_docs_parses_structured_properties(
    source: DataHubContextDocProposalsSource,
) -> None:
    alice_expert = _expert_json("urn:li:corpuser:alice")
    fin_domain = _domain_json("urn:li:domain:finance")

    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.return_value = _make_scroll_response(
        [
            {
                "urn": "urn:li:document:1",
                "experts": [alice_expert],
                "domains": [fin_domain],
            },
            {"urn": "urn:li:document:2"},
        ]
    )

    docs = source._scroll_unpublished_docs()

    assert docs[0].expert_urns == ["urn:li:corpuser:alice"]
    assert docs[0].domain_urns == ["urn:li:domain:finance"]
    assert docs[1].expert_urns == []
    assert docs[1].domain_urns == []


def test_scroll_unpublished_docs_pagination(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.side_effect = [
        _make_scroll_response([{"urn": "urn:li:document:1"}], next_scroll_id="page2"),
        _make_scroll_response([{"urn": "urn:li:document:2"}]),
    ]

    docs = source._scroll_unpublished_docs()

    assert [d.urn for d in docs] == ["urn:li:document:1", "urn:li:document:2"]
    assert mock_graph.execute_graphql.call_count == 2


def test_scroll_unpublished_docs_empty(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.return_value = _make_scroll_response([])

    docs = source._scroll_unpublished_docs()

    assert docs == []
    assert source.report.unpublished_docs_found == 0


def test_scroll_unpublished_docs_graphql_error(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.side_effect = Exception("connection refused")

    docs = source._scroll_unpublished_docs()

    assert docs == []
    assert len(source.report.failures) == 1


# ---------------------------------------------------------------------------
# _get_covered_doc_urns
# ---------------------------------------------------------------------------


def _make_action_request_info(
    created_by: str = SYSTEM_URN,
    doc_urns: Optional[List[str]] = None,
) -> ActionRequestInfoClass:
    doc_urns = doc_urns or []
    return ActionRequestInfoClass(
        type="DOCUMENTS_PROPOSAL",
        assignedUsers=[],
        assignedGroups=[],
        created=0,
        createdBy=created_by,
        params=ActionRequestParamsClass(
            documentsProposal=DocumentsProposalClass(
                proposalType=DocumentProposalTypeClass.STATE_CHANGE,
                documents=[DocumentChangeClass(targetDocument=u) for u in doc_urns],
            )
        ),
    )


def test_get_covered_doc_urns_no_existing_proposals(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.return_value = {
        "listActionRequests": {"nextScrollId": None, "actionRequests": []}
    }

    covered = source._get_covered_doc_urns()

    assert covered == set()


def test_get_covered_doc_urns_skips_non_system_proposals(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.return_value = {
        "listActionRequests": {
            "nextScrollId": None,
            "actionRequests": [{"urn": "urn:li:actionRequest:abc"}],
        }
    }
    mock_graph.get_aspect_v2.return_value = _make_action_request_info(
        created_by="urn:li:corpuser:someuser",
        doc_urns=["urn:li:document:1"],
    )

    covered = source._get_covered_doc_urns()

    assert covered == set()


def test_get_covered_doc_urns_returns_system_proposal_docs(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.return_value = {
        "listActionRequests": {
            "nextScrollId": None,
            "actionRequests": [{"urn": "urn:li:actionRequest:abc"}],
        }
    }
    mock_graph.get_aspect_v2.return_value = _make_action_request_info(
        created_by=SYSTEM_URN,
        doc_urns=["urn:li:document:1", "urn:li:document:2"],
    )

    covered = source._get_covered_doc_urns()

    assert covered == {"urn:li:document:1", "urn:li:document:2"}


def test_get_covered_doc_urns_handles_missing_aspect(
    source: DataHubContextDocProposalsSource,
) -> None:
    mock_graph = cast(MagicMock, source.graph)
    mock_graph.execute_graphql.return_value = {
        "listActionRequests": {
            "nextScrollId": None,
            "actionRequests": [{"urn": "urn:li:actionRequest:abc"}],
        }
    }
    mock_graph.get_aspect_v2.return_value = None

    covered = source._get_covered_doc_urns()

    assert covered == set()


# ---------------------------------------------------------------------------
# get_workunits (end-to-end)
# ---------------------------------------------------------------------------


def _make_doc_infos(
    urns: List[str],
    expert_urns: Optional[List[str]] = None,
    domain_urns: Optional[List[str]] = None,
) -> List[DocInfo]:
    return [
        DocInfo(
            urn=urn,
            expert_urns=list(expert_urns or []),
            domain_urns=list(domain_urns or []),
        )
        for urn in urns
    ]


def test_get_workunits_no_docs(
    source: DataHubContextDocProposalsSource,
) -> None:
    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=[]),
    ):
        workunits = list(source.get_workunits())

    assert workunits == []
    assert source.report.proposals_created == 0


def test_get_workunits_creates_proposals_for_uncovered_docs(
    source: DataHubContextDocProposalsSource,
) -> None:
    docs = _make_doc_infos(
        [f"urn:li:document:{i}" for i in range(3)],
        expert_urns=["urn:li:corpuser:alice"],
    )

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=docs),
    ):
        workunits = list(source.get_workunits())

    # One proposal → 2 workunits (actionRequestInfo + actionRequestStatus)
    assert len(workunits) == 2
    assert source.report.proposals_created == 1


def test_get_workunits_skips_covered_docs(
    source: DataHubContextDocProposalsSource,
) -> None:
    docs = _make_doc_infos(["urn:li:document:1", "urn:li:document:2"])
    covered = {"urn:li:document:1"}

    def fake_scroll() -> List[DocInfo]:
        source.report.unpublished_docs_found = len(docs)
        return docs

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=covered),
        patch.object(source, "_scroll_unpublished_docs", side_effect=fake_scroll),
    ):
        workunits = list(source.get_workunits())

    assert len(workunits) == 2  # one proposal for document:2 only
    assert source.report.proposals_created == 1
    assert source.report.docs_already_covered == 1


def test_get_workunits_batches_docs(
    source: DataHubContextDocProposalsSource,
) -> None:
    # 7 docs with same key -> 2 proposals (5 + 2)
    docs = _make_doc_infos(
        [f"urn:li:document:{i}" for i in range(7)],
        expert_urns=["urn:li:corpuser:alice"],
    )

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=docs),
    ):
        workunits = list(source.get_workunits())

    assert source.report.proposals_created == 2
    assert len(workunits) == 4  # 2 proposals × 2 MCPs each


def test_get_workunits_groups_by_experts_and_domains(
    source: DataHubContextDocProposalsSource,
) -> None:
    docs = [
        DocInfo(
            urn="urn:li:document:1",
            expert_urns=["urn:li:corpuser:alice"],
            domain_urns=["urn:li:domain:fin"],
        ),
        DocInfo(
            urn="urn:li:document:2",
            expert_urns=["urn:li:corpuser:alice"],
            domain_urns=["urn:li:domain:fin"],
        ),
        DocInfo(
            urn="urn:li:document:3",
            expert_urns=["urn:li:corpuser:bob"],
        ),
    ]

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=docs),
    ):
        workunits = list(source.get_workunits())

    # 2 groups -> 2 proposals -> 4 workunits
    assert source.report.proposals_created == 2
    assert len(workunits) == 4


def test_get_workunits_admin_role_always_assigned(
    source: DataHubContextDocProposalsSource,
) -> None:
    docs = _make_doc_infos(["urn:li:document:1"])

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=docs),
    ):
        workunits = list(source.get_workunits())

    mcp = workunits[0].metadata
    assert isinstance(mcp, MetadataChangeProposalWrapper)
    info = mcp.aspect
    assert isinstance(info, ActionRequestInfoClass)
    assert info.assignedRoles == [ADMIN_ROLE_URN]
    # No experts -> assignedUsers and assignedGroups are empty
    assert info.assignedUsers == []
    assert info.assignedGroups == []


def test_get_workunits_experts_set_as_assigned_users(
    source: DataHubContextDocProposalsSource,
) -> None:
    docs = _make_doc_infos(
        ["urn:li:document:1"],
        expert_urns=["urn:li:corpuser:alice", "urn:li:corpuser:bob"],
    )

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=docs),
    ):
        workunits = list(source.get_workunits())

    mcp = workunits[0].metadata
    assert isinstance(mcp, MetadataChangeProposalWrapper)
    info = mcp.aspect
    assert isinstance(info, ActionRequestInfoClass)
    assert set(info.assignedUsers) == {
        "urn:li:corpuser:alice",
        "urn:li:corpuser:bob",
    }
    assert info.assignedGroups == []
    # Experts present -> no admin role fallback
    assert info.assignedRoles == []


def test_get_workunits_group_experts_set_as_assigned_groups(
    source: DataHubContextDocProposalsSource,
) -> None:
    docs = _make_doc_infos(
        ["urn:li:document:1"],
        expert_urns=[
            "urn:li:corpuser:alice",
            "urn:li:corpGroup:eng",
            "urn:li:corpGroup:data",
        ],
    )

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=docs),
    ):
        workunits = list(source.get_workunits())

    mcp = workunits[0].metadata
    assert isinstance(mcp, MetadataChangeProposalWrapper)
    info = mcp.aspect
    assert isinstance(info, ActionRequestInfoClass)
    assert info.assignedUsers == ["urn:li:corpuser:alice"]
    assert set(info.assignedGroups) == {
        "urn:li:corpGroup:eng",
        "urn:li:corpGroup:data",
    }
    # Experts present -> no admin role fallback
    assert info.assignedRoles == []


def test_get_workunits_proposal_content(
    source: DataHubContextDocProposalsSource,
) -> None:
    docs = _make_doc_infos(["urn:li:document:1", "urn:li:document:2"])

    with (
        patch.object(source, "_get_covered_doc_urns", return_value=set()),
        patch.object(source, "_scroll_unpublished_docs", return_value=docs),
    ):
        workunits = list(source.get_workunits())

    assert len(workunits) == 2
    info_wu, _status_wu = workunits

    mcp = info_wu.metadata
    assert isinstance(mcp, MetadataChangeProposalWrapper)
    info = mcp.aspect
    assert isinstance(info, ActionRequestInfoClass)
    assert info.type == "DOCUMENTS_PROPOSAL"
    assert info.createdBy == SYSTEM_URN
    assert info.origin == "INFERRED"

    assert info.params is not None
    docs_proposal = info.params.documentsProposal
    assert docs_proposal is not None
    assert docs_proposal.proposalType == DocumentProposalTypeClass.STATE_CHANGE
    assert len(docs_proposal.documents) == 2
    assert all(
        d.targetLifecycleStage == PUBLISHED_STAGE_URN for d in docs_proposal.documents
    )
    target_docs = {d.targetDocument for d in docs_proposal.documents}
    assert target_docs == {"urn:li:document:1", "urn:li:document:2"}


def test_max_docs_per_proposal_constant() -> None:
    assert MAX_DOCS_PER_PROPOSAL == 5
