marker / tests /services /test_service_init.py
Vik Paruchuri
Test fix
5e316bb
import pytest
from marker.converters.pdf import PdfConverter
from marker.services.gemini import GoogleGeminiService
from marker.services.ollama import OllamaService
from marker.services.vertex import GoogleVertexService
from marker.services.openai import OpenAIService
from marker.services.azure_openai import AzureOpenAIService
@pytest.mark.output_format("markdown")
@pytest.mark.config({"page_range": [0]})
def test_empty_llm(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is None
assert pdf_converter.llm_service is None
def test_llm_no_keys(model_dict, config):
with pytest.raises(AssertionError):
PdfConverter(artifact_dict=model_dict, config={"use_llm": True})
@pytest.mark.output_format("markdown")
@pytest.mark.config({"page_range": [0], "use_llm": True, "gemini_api_key": "test"})
def test_llm_gemini(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, GoogleGeminiService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"vertex_project_id": "test",
"llm_service": "marker.services.vertex.GoogleVertexService",
}
)
def test_llm_vertex(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, GoogleVertexService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"llm_service": "marker.services.ollama.OllamaService",
}
)
def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, OllamaService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"llm_service": "marker.services.openai.OpenAIService",
"openai_api_key": "test",
}
)
def test_llm_openai(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, OpenAIService)
@pytest.mark.output_format("markdown")
@pytest.mark.config(
{
"page_range": [0],
"use_llm": True,
"llm_service": "marker.services.azure_openai.AzureOpenAIService",
"azure_endpoint": "https://example.openai.azure.com",
"azure_api_key": "test",
"deployment_name": "test-model",
"azure_api_version": "1",
}
)
def test_llm_azure_openai(pdf_converter: PdfConverter, temp_doc):
assert pdf_converter.artifact_dict["llm_service"] is not None
assert isinstance(pdf_converter.llm_service, AzureOpenAIService)