peppermenta commited on
Commit
285f1d6
·
1 Parent(s): 8c63b8f
tests/services/test_service_init.py CHANGED
@@ -4,6 +4,7 @@ from marker.converters.pdf import PdfConverter
4
  from marker.services.gemini import GoogleGeminiService
5
  from marker.services.ollama import OllamaService
6
  from marker.services.vertex import GoogleVertexService
 
7
 
8
 
9
  @pytest.mark.output_format("markdown")
@@ -38,4 +39,10 @@ def test_llm_vertex(pdf_converter: PdfConverter, temp_doc):
38
  @pytest.mark.config({"page_range": [0], "use_llm": True, "llm_service": "marker.services.ollama.OllamaService"})
39
  def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
40
  assert pdf_converter.artifact_dict["llm_service"] is not None
41
- assert isinstance(pdf_converter.llm_service, OllamaService)
 
 
 
 
 
 
 
4
  from marker.services.gemini import GoogleGeminiService
5
  from marker.services.ollama import OllamaService
6
  from marker.services.vertex import GoogleVertexService
7
+ from marker.services.openai import OpenAIService
8
 
9
 
10
  @pytest.mark.output_format("markdown")
 
39
  @pytest.mark.config({"page_range": [0], "use_llm": True, "llm_service": "marker.services.ollama.OllamaService"})
40
  def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
41
  assert pdf_converter.artifact_dict["llm_service"] is not None
42
+ assert isinstance(pdf_converter.llm_service, OllamaService)
43
+
44
+ @pytest.mark.output_format("markdown")
45
+ @pytest.mark.config({"page_range": [0], "use_llm": True, "llm_service": "marker.services.openai.OpenAIService", "openai_api_key": "test"})
46
+ def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
47
+ assert pdf_converter.artifact_dict["llm_service"] is not None
48
+ assert isinstance(pdf_converter.llm_service, OpenAIService)