unstructured

1
from unstructured.documents.elements import Text
2
from unstructured.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
3

4

5
def test_embed_documents_does_not_break_element_to_dict(mocker):
6
    # Mocked client with the desired behavior for embed_documents
7
    mock_client = mocker.MagicMock()
8
    mock_client.embed_documents.return_value = [1, 2]
9

10
    # Mock create_client to return our mock_client
11
    mocker.patch.object(OpenAIEmbeddingEncoder, "create_client", return_value=mock_client)
12

13
    encoder = OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(api_key="api_key"))
14
    elements = encoder.embed_documents(
15
        elements=[Text("This is sentence 1"), Text("This is sentence 2")],
16
    )
17
    assert len(elements) == 2
18
    assert elements[0].to_dict()["text"] == "This is sentence 1"
19
    assert elements[1].to_dict()["text"] == "This is sentence 2"
20