llama-hub/tests/tests_confluence/test_confluence_reader.py

import pytest
import unittest
from unittest.mock import patch
from loader_hub.confluence.base import ConfluenceReader, Document


@pytest.fixture
def mock_confluence():
    with patch("atlassian.Confluence") as mock_confluence:
        yield mock_confluence


CONFLUENCE_BASE_URL = "https://example.atlassian.com/wiki"
MOCK_OAUTH = {
    "client_id": "your_client_id",
    "token": {
        "access_token": "your_access_token",
        "token_type": "Bearer",
    },
}


class TestConfluenceReader:
    def test_confluence_reader_initialization(self, mock_confluence):

        # Test with oauth2
        ConfluenceReader(base_url=CONFLUENCE_BASE_URL, oauth2=MOCK_OAUTH)
        mock_confluence.assert_called_once_with(
            url=CONFLUENCE_BASE_URL, oauth2=MOCK_OAUTH, cloud=True
        )

        # Test without oauth2
        with unittest.mock.patch.dict(
            "os.environ",
            {"CONFLUENCE_USERNAME": "user", "CONFLUENCE_API_TOKEN": "api_token"},
        ):
            ConfluenceReader(base_url=CONFLUENCE_BASE_URL)
            mock_confluence.assert_called_with(
                url=CONFLUENCE_BASE_URL,
                username="user",
                password="api_token",
                cloud=True,
            )

    def test_confluence_reader_load_data_invalid_args(self, mock_confluence):
        confluence_reader = ConfluenceReader(
            base_url=CONFLUENCE_BASE_URL, oauth2=MOCK_OAUTH
        )
        confluence_reader.confluence = mock_confluence

        with pytest.raises(
            ValueError,
            match="Must specify at least one among `space_key`, `page_ids`, `label`, `cql` parameters.",
        ):
            confluence_reader.load_data()

    def test_confluence_reader_load_data_by_page_ids(self, mock_confluence):
        mock_confluence.get_page_by_id.side_effect = [
            {
                "id": "123",
                "title": "Page 123",
                "body": {"storage": {"value": "<p>Content 123</p>"}},
            },
            {
                "id": "456",
                "title": "Page 456",
                "body": {"storage": {"value": "<p>Content 456</p>"}},
            },
        ]

        confluence_reader = ConfluenceReader(
            base_url=CONFLUENCE_BASE_URL, oauth2=MOCK_OAUTH
        )
        confluence_reader.confluence = mock_confluence

        mock_page_ids = ["123", "456"]
        documents = confluence_reader.load_data(page_ids=mock_page_ids)

        assert len(documents) == 2
        assert all(isinstance(doc, Document) for doc in documents)
        assert documents[0].doc_id == "123"
        assert documents[0].extra_info == {"title": "Page 123"}
        assert documents[1].doc_id == "456"
        assert documents[1].extra_info == {"title": "Page 456"}

        assert mock_confluence.get_page_by_id.call_count == 2

        assert mock_confluence.get_all_pages_from_space.call_count == 0
        assert mock_confluence.get_all_pages_by_label.call_count == 0
        assert mock_confluence.cql.call_count == 0
        assert mock_confluence.get_page_child_by_type.call_count == 0

    def test_confluence_reader_load_data_by_space_id(self, mock_confluence):
        # one response with two pages
        mock_confluence.get_all_pages_from_space.return_value = [
            {
                "id": "123",
                "type": "page",
                "status": "current",
                "title": "Page 123",
                "body": {"storage": {"value": "<p>Content 123</p>"}},
            },
            {
                "id": "456",
                "type": "page",
                "status": "current",
                "title": "Page 456",
                "body": {"storage": {"value": "<p>Content 456</p>"}},
            },
        ]

        confluence_reader = ConfluenceReader(
            base_url=CONFLUENCE_BASE_URL, oauth2=MOCK_OAUTH
        )
        confluence_reader.confluence = mock_confluence

        mock_space_key = "spaceId123"
        documents = confluence_reader.load_data(space_key=mock_space_key)

        assert mock_confluence.get_all_pages_from_space.call_count == 1
        assert mock_confluence.get_all_pages_from_space.call_args[0][0] == "spaceId123"
        assert mock_confluence.get_all_pages_from_space.call_args[1]["start"] == 0
        assert mock_confluence.get_all_pages_from_space.call_args[1]["limit"] == 50

        assert len(documents) == 2
        assert all(isinstance(doc, Document) for doc in documents)
        assert documents[0].doc_id == "123"
        assert documents[0].extra_info == {"title": "Page 123"}
        assert documents[1].doc_id == "456"
        assert documents[1].extra_info == {"title": "Page 456"}

        assert mock_confluence.get_page_by_id.call_count == 0
        assert mock_confluence.get_all_pages_by_label.call_count == 0
        assert mock_confluence.cql.call_count == 0
        assert mock_confluence.get_page_child_by_type.call_count == 0

    def test_confluence_reader_load_data_by_space_id_pagination(self, mock_confluence):
        # two api responses with one page each
        mock_confluence.get_all_pages_from_space.side_effect = [
            [
                {
                    "id": "123",
                    "type": "page",
                    "status": "current",
                    "title": "Page 123",
                    "body": {"storage": {"value": "<p>Content 123</p>"}},
                },
            ],
            [
                {
                    "id": "456",
                    "type": "page",
                    "status": "current",
                    "title": "Page 456",
                    "body": {"storage": {"value": "<p>Content 456</p>"}},
                }
            ],
            [],
        ]

        confluence_reader = ConfluenceReader(
            base_url=CONFLUENCE_BASE_URL, oauth2=MOCK_OAUTH
        )
        confluence_reader.confluence = mock_confluence

        mock_space_key = "spaceId123"
        mock_limit = 1  # fetch one page at a time
        documents = confluence_reader.load_data(
            space_key=mock_space_key, limit=mock_limit
        )

        assert mock_confluence.get_all_pages_from_space.call_count == 3

        assert len(documents) == 2
        assert all(isinstance(doc, Document) for doc in documents)
        assert documents[0].doc_id == "123"
        assert documents[0].extra_info == {"title": "Page 123"}
        assert documents[1].doc_id == "456"
        assert documents[1].extra_info == {"title": "Page 456"}

        assert mock_confluence.get_page_by_id.call_count == 0
        assert mock_confluence.get_all_pages_by_label.call_count == 0
        assert mock_confluence.cql.call_count == 0
        assert mock_confluence.get_page_child_by_type.call_count == 0