#!/usr/bin/env python3

"""

Comprehensive test suite for the PDF field extraction CLI.

Tests both the extractor module and the CLI interface.

"""

import json

import os

import sys

import tempfile

import unittest

from pathlib import Path

from unittest.mock import patch, MagicMock, mock_open

import argparse

# Add the project root to Python path

sys.path.insert(0, str(Path(__file__).parent))

from extractor import DoclingExtractor

from extract import load_fields_config, validate_pdf_file, save_results, create_output_directory, main

class TestDoclingExtractor(unittest.TestCase):

"""Test cases for DoclingExtractor class."""

def setUp(self):

"""Set up test fixtures."""

self.temp_dir = tempfile.mkdtemp()

self.extractor = DoclingExtractor(output_dir=self.temp_dir)

# Sample fields configuration

self.sample_fields = {

"INVOICE NUMBER": "string",

"Date": "date",

"Total": "number",

"BILL TO address": "string"

}

# Sample document text

self.sample_text = """

INVOICE NUMBER: INV-12345

Date: 2024-01-15

BILL TO:

Test Company

123 Test Street

Test City, TS 12345

Items:

Item 1 $100.00

Item 2 $50.00

Total: $150.00

"""

def test_initialization(self):

"""Test extractor initialization."""

self.assertIsNotNone(self.extractor.converter)

self.assertIsNotNone(self.extractor.llm)

self.assertEqual(self.extractor.output_dir, Path(self.temp_dir))

self.assertTrue(self.extractor.output_dir.exists())

def test_build_extraction_prompt(self):

"""Test prompt building with different field configurations."""

# Test with dict

prompt = self.extractor.build_extraction_prompt(self.sample_fields)

self.assertIsNotNone(prompt)

# Test with list

field_list = list(self.sample_fields.keys())

prompt = self.extractor.build_extraction_prompt(field_list)

self.assertIsNotNone(prompt)

@patch('extractor.ChatVertexAI')

def test_extract_fields_mock(self, mock_llm_class):

"""Test field extraction with mocked LLM."""

# Mock LLM response

mock_llm = MagicMock()

mock_llm_class.return_value = mock_llm

mock_response = MagicMock()

mock_response.content = json.dumps({

"INVOICE NUMBER": "INV-12345",

"Date": "2024-01-15",

"Total": "$150.00",

"BILL TO address": "Test Company\n123 Test Street\nTest City, TS 12345"

})

# Re-initialize extractor with mocked LLM

extractor = DoclingExtractor(output_dir=self.temp_dir)

# Mock the chain invoke

with patch.object(extractor, '_extract_with_llm') as mock_extract:

mock_extract.return_value = {

"INVOICE NUMBER": "INV-12345",

"Date": "2024-01-15",

"Total": "$150.00",

"BILL TO address": "Test Company\n123 Test Street\nTest City, TS 12345"

}

result = extractor.extract_fields(self.sample_text, self.sample_fields)

self.assertIn("INVOICE NUMBER", result)

self.assertEqual(result["INVOICE NUMBER"], "INV-12345")

def test_check_missing_fields(self):

"""Test missing field detection."""

# All fields present

extracted = {

"INVOICE NUMBER": "INV-12345",

"Date": "2024-01-15",

"Total": "$150.00",

"BILL TO address": "Test Company"

}

missing = self.extractor.check_missing_fields(extracted, self.sample_fields)

self.assertEqual(len(missing), 0)

# Some fields missing

extracted_partial = {

"INVOICE NUMBER": "INV-12345",

"Date": None,

"Total": "",

"BILL TO address": "Test Company"

}

missing = self.extractor.check_missing_fields(extracted_partial, self.sample_fields)

self.assertEqual(len(missing), 2)

self.assertIn("Date", missing)

self.assertIn("Total", missing)

class TestCLIFunctions(unittest.TestCase):

"""Test cases for CLI helper functions."""

def setUp(self):

"""Set up test fixtures."""

self.temp_dir = tempfile.mkdtemp()

# Sample fields config

self.sample_config = {

"INVOICE NUMBER": "string",

"Date": "date",

"Total": "number"

}

def test_load_fields_config_valid(self):

"""Test loading valid fields configuration."""

config_file = Path(self.temp_dir) / "test_fields.json"

with open(config_file, 'w') as f:

json.dump(self.sample_config, f)

config = load_fields_config(str(config_file))

self.assertEqual(config, self.sample_config)

def test_load_fields_config_invalid_json(self):

"""Test loading invalid JSON configuration."""

config_file = Path(self.temp_dir) / "invalid.json"

with open(config_file, 'w') as f:

f.write("{ invalid json }")

with self.assertRaises(SystemExit):

load_fields_config(str(config_file))

def test_load_fields_config_invalid_types(self):

"""Test loading configuration with unsupported field types."""

invalid_config = {

"field1": "unsupported_type"

}

config_file = Path(self.temp_dir) / "invalid_types.json"

with open(config_file, 'w') as f:

json.dump(invalid_config, f)

with self.assertRaises(SystemExit):

load_fields_config(str(config_file))

def test_load_fields_config_not_dict(self):

"""Test loading configuration that's not a dict."""

invalid_config = ["field1", "field2"]

config_file = Path(self.temp_dir) / "invalid_format.json"

with open(config_file, 'w') as f:

json.dump(invalid_config, f)

with self.assertRaises(SystemExit):

load_fields_config(str(config_file))

def test_validate_pdf_file_valid(self):

"""Test PDF file validation with valid file."""

# Create a dummy PDF file

pdf_file = Path(self.temp_dir) / "test.pdf"

pdf_file.write_bytes(b"%PDF-1.4 dummy content")

result = validate_pdf_file(str(pdf_file))

self.assertEqual(result, pdf_file)

def test_validate_pdf_file_not_found(self):

"""Test PDF file validation with non-existent file."""

with self.assertRaises(SystemExit):

validate_pdf_file("nonexistent.pdf")

def test_validate_pdf_file_wrong_extension(self):

"""Test PDF file validation with wrong file extension."""

txt_file = Path(self.temp_dir) / "test.txt"

txt_file.write_text("Not a PDF")

with self.assertRaises(SystemExit):

validate_pdf_file(str(txt_file))

def test_validate_pdf_file_not_file(self):

"""Test PDF file validation with directory."""

with self.assertRaises(SystemExit):

validate_pdf_file(str(self.temp_dir))

def test_create_output_directory(self):

"""Test output directory creation."""

output_dir = create_output_directory()

self.assertTrue(output_dir.exists())

self.assertEqual(output_dir.name, "output")

def test_save_results(self):

"""Test saving extraction results."""

results = {

"success": True,

"extracted_fields": {"INVOICE NUMBER": "INV-123"},

"missing_fields": []

}

output_file = save_results(results, "test_pdf", Path(self.temp_dir))

self.assertTrue(output_file.exists())

self.assertEqual(output_file.name, "extracted_test_pdf.json")

# Verify content

with open(output_file, 'r') as f:

saved_data = json.load(f)

self.assertEqual(saved_data, results)

@patch('extract.sys.argv', ['extract.py', 'fields.json', 'test.pdf'])

@patch('extract.load_fields_config')

@patch('extract.validate_pdf_file')

@patch('extract.create_output_directory')

@patch('extract.DoclingExtractor')

def test_main_success(self, mock_extractor_class, mock_create_dir, mock_validate, mock_load_config):

"""Test main CLI function with successful extraction."""

# Setup mocks

mock_load_config.return_value = {"field1": "string"}

mock_validate.return_value = Path("test.pdf")

mock_create_dir.return_value = Path("output")

mock_extractor = MagicMock()

mock_extractor.process_pdf_with_logging.return_value = {

"success": True,

"extracted_fields": {"field1": "value1"},

"missing_fields": []

}

mock_extractor_class.return_value = mock_extractor

# Mock save_results to avoid file operations

with patch('extract.save_results') as mock_save:

mock_save.return_value = Path("output/result.json")

# Test main function

with patch('extract.Path') as mock_path:

mock_path.return_value = Path("test.pdf")

main()

# Verify calls

mock_load_config.assert_called_once_with('fields.json')

mock_validate.assert_called_once_with('test.pdf')

mock_extractor.process_pdf_with_logging.assert_called_once()

@patch('extract.sys.argv', ['extract.py', 'fields.json', 'test.pdf'])

@patch('extract.load_fields_config')

def test_main_load_config_error(self, mock_load_config):

"""Test main function with config loading error."""

mock_load_config.side_effect = SystemExit(1)

with self.assertRaises(SystemExit):

main()

@patch('extract.sys.argv', ['extract.py', 'fields.json', 'test.pdf'])

@patch('extract.load_fields_config')

@patch('extract.validate_pdf_file')

@patch('extract.DoclingExtractor')

def test_main_extraction_error(self, mock_extractor_class, mock_validate, mock_load_config):

"""Test main function with extraction error."""

# Setup mocks

mock_load_config.return_value = {"field1": "string"}

mock_validate.return_value = Path("test.pdf")

mock_extractor = MagicMock()

mock_extractor.process_pdf_with_logging.side_effect = Exception("Extraction failed")

mock_extractor_class.return_value = mock_extractor

with self.assertRaises(SystemExit):

main()

class TestAdditionalExtractorMethods(unittest.TestCase):

"""Test additional extractor methods for better coverage."""

def setUp(self):

"""Set up test fixtures."""

self.temp_dir = tempfile.mkdtemp()

self.extractor = DoclingExtractor(output_dir=self.temp_dir)

def test_save_raw_output(self):

"""Test saving raw output files."""

raw_output = {

"markdown": "# Test Document\nContent here",

"metadata": {

"num_pages": 1,

"source": "test.pdf",

"parser": "docling"

}

}

result_path = self.extractor.save_raw_output(raw_output, 1, "test_pdf")

self.assertTrue(Path(result_path).exists())

self.assertTrue(result_path.endswith('.json'))

# Check JSON file was created

with open(result_path, 'r') as f:

saved_data = json.load(f)

self.assertEqual(saved_data, raw_output)

# Check markdown file was also created

md_file = result_path.replace('.json', '.md')

self.assertTrue(Path(md_file).exists())

def test_save_audit_log(self):

"""Test saving audit log."""

result = {

"success": True,

"attempts": 1,

"extracted_fields": {"field1": "value1"},

"missing_fields": [],

"raw_output_files": ["test.json"],

"error": None,

"logs": [{"level": "info", "message": "test"}]

}

fields = {"field1": "string"}

audit_path = self.extractor.save_audit_log("test_pdf", result, fields)

self.assertTrue(Path(audit_path).exists())

self.assertTrue(audit_path.endswith('.json'))

# Verify audit content

with open(audit_path, 'r') as f:

audit_data = json.load(f)

self.assertIn("extraction_summary", audit_data)

self.assertIn("field_analysis", audit_data)

self.assertEqual(audit_data["extraction_summary"]["pdf_name"], "test_pdf")

def test_check_missing_fields_dict(self):

"""Test missing field check with dict input."""

fields = {"field1": "string", "field2": "number"}

extracted = {"field1": "value1", "field2": None}

missing = self.extractor.check_missing_fields(extracted, fields)

self.assertEqual(missing, ["field2"])

def test_check_missing_fields_list(self):

"""Test missing field check with list input."""

fields = ["field1", "field2", "field3"]

extracted = {"field1": "value1", "field2": "", "field3": "value3"}

missing = self.extractor.check_missing_fields(extracted, fields)

self.assertEqual(missing, ["field2"])

def test_check_missing_fields_null_values(self):

"""Test missing field check with null string values."""

fields = ["field1", "field2"]

extracted = {"field1": "value1", "field2": "null"}

missing = self.extractor.check_missing_fields(extracted, fields)

self.assertEqual(missing, ["field2"])

def test_build_extraction_prompt_dict(self):

"""Test prompt building with dict fields."""

fields = {"field1": "string", "field2": "number"}

prompt = self.extractor.build_extraction_prompt(fields)

self.assertIsNotNone(prompt)

# Check that prompt contains field information

prompt_str = str(prompt)

self.assertIn("field1", prompt_str)

self.assertIn("field2", prompt_str)

def test_build_extraction_prompt_list(self):

"""Test prompt building with list fields."""

fields = ["field1", "field2"]

prompt = self.extractor.build_extraction_prompt(fields)

self.assertIsNotNone(prompt)

prompt_str = str(prompt)

self.assertIn("field1", prompt_str)

self.assertIn("field2", prompt_str)

def test_parse_pdf(self):

"""Test PDF parsing method."""

# Mock the DocumentConverter

with patch.object(self.extractor.converter, 'convert') as mock_convert:

mock_result = MagicMock()

mock_result.document.export_to_markdown.return_value = "# Test Document\nContent"

mock_result.document.pages = [MagicMock()] # Mock 1 page

mock_convert.return_value = mock_result

text, raw_output = self.extractor.parse_pdf("test.pdf")

self.assertEqual(text, "# Test Document\nContent")

self.assertIn("markdown", raw_output)

self.assertIn("metadata", raw_output)

self.assertEqual(raw_output["metadata"]["source"], "test.pdf")

@patch('extractor.px.launch_app')

@patch('extractor.register')

@patch('extractor.LangChainInstrumentor')

def test_setup_phoenix_tracing_success(self, mock_instrumentor, mock_register, mock_launch):

"""Test successful Phoenix tracing setup."""

# Reset the class variable to test setup

DoclingExtractor._phoenix_initialized = False

# Enable Phoenix for this test

with patch.dict(os.environ, {'PHOENIX_ENABLED': 'true'}):

# Mock the register and instrumentor

mock_tracer = MagicMock()

mock_register.return_value = mock_tracer

# Call setup method

DoclingExtractor._setup_phoenix_tracing()

# Verify calls

mock_register.assert_called_once()

mock_instrumentor.return_value.instrument.assert_called_once_with(tracer_provider=mock_tracer)

mock_launch.assert_called_once()

self.assertTrue(DoclingExtractor._phoenix_initialized)

@patch('extractor.px.launch_app')

@patch('extractor.register')

def test_setup_phoenix_tracing_failure(self, mock_register, mock_launch):

"""Test Phoenix tracing setup failure."""

# Reset the class variable

DoclingExtractor._phoenix_initialized = False

# Enable Phoenix for this test

with patch.dict(os.environ, {'PHOENIX_ENABLED': 'true'}):

# Mock failure

mock_register.side_effect = Exception("Phoenix setup failed")

# Call setup method (should not raise exception)

DoclingExtractor._setup_phoenix_tracing()

# Should not be initialized

self.assertFalse(DoclingExtractor._phoenix_initialized)

def test_setup_phoenix_tracing_disabled(self):

"""Test Phoenix tracing setup when disabled."""

# Reset the class variable

DoclingExtractor._phoenix_initialized = False

# Ensure Phoenix is disabled

with patch.dict(os.environ, {'PHOENIX_ENABLED': 'false'}):

# Call setup method

DoclingExtractor._setup_phoenix_tracing()

# Should be marked as initialized to avoid repeated checks

self.assertTrue(DoclingExtractor._phoenix_initialized)

def test_process_pdf_with_custom_fields(self):

"""Test process_pdf with custom fields."""

# Mock the workflow

with patch.object(self.extractor.workflow, 'invoke') as mock_invoke:

mock_invoke.return_value = {

"success": True,

"attempt": 1,

"extracted_fields": {"field1": "value1"},

"missing_fields": [],

"raw_output_files": [],

"error": None,

"logs": []

}

result = self.extractor.process_pdf("test.pdf", custom_fields=["field1"])

self.assertTrue(result["success"])

self.assertIn("field1", result["extracted_fields"])

def test_process_pdf_without_custom_fields(self):

"""Test process_pdf with default fields."""

# Mock the workflow

with patch.object(self.extractor.workflow, 'invoke') as mock_invoke:

mock_invoke.return_value = {

"success": True,

"attempt": 1,

"extracted_fields": {"Shipping Charges": "100"},

"missing_fields": [],

"raw_output_files": [],

"error": None,

"logs": []

}

result = self.extractor.process_pdf("test.pdf")

self.assertTrue(result["success"])

mock_invoke.assert_called_once()

def test_process_pdf_streaming(self):

"""Test process_pdf_streaming method."""

# Mock the workflow

with patch.object(self.extractor.workflow, 'invoke') as mock_invoke:

mock_invoke.return_value = {

"success": True,

"attempt": 1,

"extracted_fields": {"field1": "value1"},

"missing_fields": [],

"raw_output_files": [],

"error": None,

"logs": []

}

# Test with callback

callback_calls = []

def test_callback(msg):

callback_calls.append(msg)

result = self.extractor.process_pdf_streaming("test.pdf", custom_fields=["field1"], log_callback=test_callback)

self.assertTrue(result["success"])

self.assertIn("field1", result["extracted_fields"])

def test_extract_fields_with_chain(self):

"""Test extract_fields method with chain."""

# Mock the chain

mock_chain = MagicMock()

mock_response = MagicMock()

mock_response.content = '{"field1": "value1"}'

mock_chain.invoke.return_value = mock_response

result = self.extractor.extract_fields("Sample text", {"field1": "string"}, mock_chain)

# The method might return None values, so just check it doesn't crash

self.assertIsNotNone(result)

def test_extract_fields_json_error(self):

"""Test extract_fields method with JSON error."""

# Mock the chain with invalid JSON

mock_chain = MagicMock()

mock_response = MagicMock()

mock_response.content = '{ invalid json }'

mock_chain.invoke.return_value = mock_response

result = self.extractor.extract_fields("Sample text", {"field1": "string"}, mock_chain)

# The method might return None values, so just check it doesn't crash

self.assertIsNotNone(result)

def test_extract_fields_chain_error(self):

"""Test extract_fields method with chain error."""

# Mock the chain with error

mock_chain = MagicMock()

mock_chain.invoke.side_effect = Exception("Chain error")

result = self.extractor.extract_fields("Sample text", {"field1": "string"}, mock_chain)

# Should return empty dict or None values on error

self.assertIsInstance(result, dict)

class TestPhoenixIntegration(unittest.TestCase):

"""Test cases for Arize Phoenix integration."""

def test_phoenix_config_variables(self):

"""Test Phoenix configuration environment variables."""

# Check default values (these may be overridden by environment)

self.assertIsInstance(DoclingExtractor.PHOENIX_PROJECT_NAME, str)

self.assertEqual(DoclingExtractor.PHOENIX_COLLECTOR_ENDPOINT, "http://localhost:6006")

@patch.dict(os.environ, {'PHOENIX_PROJECT_NAME': 'test-project'})

def test_phoenix_env_override(self):

"""Test Phoenix configuration via environment variables."""

# This should pick up the environment variable

extractor = DoclingExtractor()

self.assertIsNotNone(extractor)

class TestWorkflowMethods(unittest.TestCase):

"""Test workflow methods for better coverage."""

def setUp(self):

"""Set up test fixtures."""

self.temp_dir = tempfile.mkdtemp()

self.extractor = DoclingExtractor(output_dir=self.temp_dir)

def test_parse_pdf_node(self):

"""Test the parse_pdf workflow node."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"field_names": ["field1"],

"attempt": 1,

"logs": []

}

# Mock the parse_pdf method

with patch.object(self.extractor, 'parse_pdf') as mock_parse:

mock_parse.return_value = ("Sample text", {"markdown": "Sample", "metadata": {}})

result = self.extractor._parse_pdf_node(state)

self.assertEqual(result["text_content"], "Sample text")

self.assertIn("raw_output", result)

def test_extract_fields_node(self):

"""Test the extract_fields workflow node."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"field_names": ["field1"],

"text_content": "Sample text with field1: value1",

"attempt": 1,

"raw_output_files": [],

"logs": [],

"raw_output": {"markdown": "Sample", "metadata": {}}

}

# Mock the extract_fields method

with patch.object(self.extractor, 'extract_fields') as mock_extract:

mock_extract.return_value = {"field1": "value1"}

result = self.extractor._extract_fields_node(state)

self.assertIn("extracted_fields", result)

self.assertEqual(result["extracted_fields"]["field1"], "value1")

def test_validate_extraction_node_success(self):

"""Test the validate_extraction workflow node with success."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string", "field2": "number"},

"field_names": ["field1", "field2"],

"extracted_fields": {"field1": "value1", "field2": "100"},

"attempt": 1,

"raw_output_files": ["test.json"],

"logs": []

}

result = self.extractor._validate_extraction_node(state)

self.assertTrue(result.get("success", False))

self.assertEqual(result.get("missing_fields", []), [])

def test_validate_extraction_node_missing_fields(self):

"""Test the validate_extraction workflow node with missing fields."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string", "field2": "number"},

"field_names": ["field1", "field2"],

"extracted_fields": {"field1": "value1"},

"attempt": 1,

"raw_output_files": ["test.json"],

"logs": []

}

result = self.extractor._validate_extraction_node(state)

# Check if missing fields are detected

missing = result.get("missing_fields", [])

self.assertIn("field2", missing)

def test_validate_extraction_node_retry_limit(self):

"""Test the validate_extraction workflow node at retry limit."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"field_names": ["field1"],

"extracted_fields": {},

"attempt": 5, # At retry limit

"raw_output_files": ["test.json"],

"logs": []

}

result = self.extractor._validate_extraction_node(state)

# At retry limit, should mark as success even with missing fields

self.assertTrue(result.get("success", False))

self.assertIn("missing_fields", result)

def test_should_retry_condition(self):

"""Test the should_retry workflow condition."""

# Test when should retry

state_false = {

"success": False,

"attempt": 1,

"missing_fields": ["field1"]

}

retry_result = self.extractor._should_retry(state_false)

self.assertFalse(retry_result) # Should return False for retry

# Test when should not retry (success)

state_true = {

"success": True,

"attempt": 1,

"missing_fields": []

}

success_result = self.extractor._should_retry(state_true)

self.assertTrue(success_result) # Should return True for END

# Test when should not retry (at limit)

state_limit = {

"success": False,

"attempt": 5,

"missing_fields": ["field1"]

}

limit_result = self.extractor._should_retry(state_limit)

self.assertTrue(limit_result) # Should return True for END

def test_extract_with_llm_success(self):

"""Test _extract_with_llm method with success."""

text = "Sample text with field1: value1"

fields = {"field1": "string"}

# Mock the chain

mock_chain = MagicMock()

mock_response = MagicMock()

mock_response.content = '{"field1": "value1"}'

mock_chain.invoke.return_value = mock_response

result = self.extractor._extract_with_llm(text, fields, mock_chain)

self.assertEqual(result["field1"], "value1")

def test_extract_with_llm_json_error(self):

"""Test _extract_with_llm method with JSON parsing error."""

text = "Sample text"

fields = {"field1": "string"}

# Mock the chain with invalid JSON

mock_chain = MagicMock()

mock_response = MagicMock()

mock_response.content = '{ invalid json }'

mock_chain.invoke.return_value = mock_response

result = self.extractor._extract_with_llm(text, fields, mock_chain)

self.assertEqual(result, {})

def test_extract_with_llm_chain_error(self):

"""Test _extract_with_llm method with chain error."""

text = "Sample text"

fields = {"field1": "string"}

# Mock the chain with error

mock_chain = MagicMock()

mock_chain.invoke.side_effect = Exception("Chain error")

result = self.extractor._extract_with_llm(text, fields, mock_chain)

self.assertEqual(result, {})

def test_process_pdf_with_logging(self):

"""Test process_pdf_with_logging method."""

# Mock the workflow

with patch.object(self.extractor.workflow, 'invoke') as mock_invoke:

mock_invoke.return_value = {

"success": True,

"attempt": 1,

"extracted_fields": {"field1": "value1"},

"missing_fields": [],

"raw_output_files": [],

"error": None,

"logs": []

}

result = self.extractor.process_pdf_with_logging("test.pdf", custom_fields=["field1"])

self.assertTrue(result["success"])

self.assertIn("field1", result["extracted_fields"])

class TestEdgeCases(unittest.TestCase):

"""Test edge cases and error handling."""

def setUp(self):

"""Set up test fixtures."""

self.temp_dir = tempfile.mkdtemp()

self.extractor = DoclingExtractor(output_dir=self.temp_dir)

def test_build_extraction_prompt_empty_dict(self):

"""Test prompt building with empty dict."""

prompt = self.extractor.build_extraction_prompt({})

self.assertIsNotNone(prompt)

def test_build_extraction_prompt_empty_list(self):

"""Test prompt building with empty list."""

prompt = self.extractor.build_extraction_prompt([])

self.assertIsNotNone(prompt)

def test_check_missing_fields_empty_extracted(self):

"""Test missing field check with empty extracted dict."""

fields = ["field1", "field2"]

extracted = {}

missing = self.extractor.check_missing_fields(extracted, fields)

self.assertEqual(set(missing), {"field1", "field2"})

def test_check_missing_fields_all_present(self):

"""Test missing field check with all fields present."""

fields = ["field1", "field2"]

extracted = {"field1": "value1", "field2": "value2"}

missing = self.extractor.check_missing_fields(extracted, fields)

self.assertEqual(missing, [])

def test_save_raw_output_with_metadata(self):

"""Test saving raw output with complete metadata."""

raw_output = {

"markdown": "# Test\nContent",

"metadata": {

"num_pages": 2,

"source": "test.pdf",

"parser": "docling",

"processing_time": 1.5

}

}

result_path = self.extractor.save_raw_output(raw_output, 1, "test_pdf")

self.assertTrue(Path(result_path).exists())

with open(result_path, 'r') as f:

saved_data = json.load(f)

self.assertEqual(saved_data["metadata"]["num_pages"], 2)

self.assertEqual(saved_data["metadata"]["processing_time"], 1.5)

class TestWorkflowComprehensive(unittest.TestCase):

"""Comprehensive workflow tests to reach 85% coverage."""

def setUp(self):

"""Set up test fixtures."""

self.temp_dir = tempfile.mkdtemp()

self.extractor = DoclingExtractor(output_dir=self.temp_dir)

def test_parse_pdf_node_complete(self):

"""Test the parse_pdf workflow node with complete state."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"field_names": ["field1"],

"attempt": 1,

"logs": []

}

# Mock the parse_pdf method

with patch.object(self.extractor, 'parse_pdf') as mock_parse:

mock_parse.return_value = ("Sample text", {"markdown": "Sample", "metadata": {"source": "test.pdf"}})

result = self.extractor._parse_pdf_node(state)

self.assertEqual(result["text_content"], "Sample text")

self.assertIn("raw_output", result)

self.assertIn("logs", result)

def test_extract_fields_node_complete(self):

"""Test the extract_fields workflow node with complete state."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"field_names": ["field1"],

"text_content": "Sample text",

"attempt": 1,

"raw_output_files": [],

"logs": [],

"raw_output": {"markdown": "Sample", "metadata": {"source": "test.pdf"}}

}

# Mock the extract_fields method

with patch.object(self.extractor, 'extract_fields') as mock_extract:

mock_extract.return_value = {"field1": "value1"}

# Mock save_raw_output

with patch.object(self.extractor, 'save_raw_output') as mock_save:

mock_save.return_value = "output/raw.json"

result = self.extractor._extract_fields_node(state)

self.assertIn("extracted_fields", result)

self.assertIn("logs", result)

def test_validate_extraction_node_complete(self):

"""Test the validate_extraction workflow node with complete scenarios."""

# Test success case

state_success = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"field_names": ["field1"],

"extracted_fields": {"field1": "value1"},

"attempt": 1,

"raw_output_files": ["test.json"],

"logs": []

}

result = self.extractor._validate_extraction_node(state_success)

self.assertTrue(result.get("success", False))

# Test missing fields case

state_missing = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"field_names": ["field1"],

"extracted_fields": {},

"attempt": 1,

"raw_output_files": ["test.json"],

"logs": []

}

result = self.extractor._validate_extraction_node(state_missing)

self.assertFalse(result.get("success", True))

self.assertIn("missing_fields", result)

def test_should_retry_conditions(self):

"""Test all should_retry workflow conditions."""

# Test retry case

state_retry = {

"success": False,

"attempt": 1,

"missing_fields": ["field1"]

}

result = self.extractor._should_retry(state_retry)

self.assertFalse(result) # Should return "retry"

# Test success case

state_success = {

"success": True,

"attempt": 1,

"missing_fields": []

}

result = self.extractor._should_retry(state_success)

self.assertTrue(result) # Should return END

# Test limit case

state_limit = {

"success": False,

"attempt": 5,

"missing_fields": ["field1"]

}

result = self.extractor._should_retry(state_limit)

self.assertTrue(result) # Should return END

def test_save_results_node_complete(self):

"""Test the save_results workflow node."""

state = {

"pdf_path": "test.pdf",

"pdf_name": "test.pdf",

"fields_to_extract": {"field1": "string"},

"success": True,

"attempt": 1,

"extracted_fields": {"field1": "value1"},

"missing_fields": [],

"raw_output_files": ["test.json"],

"error": None,

"logs": []

}

# Mock save methods

with patch('extract.save_results') as mock_save:

with patch.object(self.extractor, 'save_audit_log') as mock_audit:

mock_save.return_value = Path("output/result.json")

mock_audit.return_value = Path("output/audit.json")

result = self.extractor._save_results_node(state)

self.assertTrue(result.get("success", False))

def test_extract_with_llm_complete(self):

"""Test _extract_with_llm method with all scenarios."""

text = "Sample text"

fields = {"field1": "string"}

# Test success

mock_chain = MagicMock()