"""
Unit tests for HtmlTableExtractor.
Tests the HTML table extraction functionality including:
- Table extraction from HTML
- Classification of valid data tables vs raw text tables
- Row and column handling
"""
import pytest
from src.extraction_v2.html_table_extractor import (
HtmlTableExtractor,
HtmlTable,
extract_tables_from_html,
extract_all_tables_from_html,
)
class TestHtmlTableExtractorInit:
"""Test HtmlTableExtractor initialization."""
def test_default_min_cols(self):
"""Test default minimum columns threshold."""
extractor = HtmlTableExtractor()
assert extractor.MIN_TABLE_COLS == 2
def test_custom_min_cols(self):
"""Test custom minimum columns threshold."""
extractor = HtmlTableExtractor(min_cols=3)
assert extractor.MIN_TABLE_COLS == 3
class TestExtractAllTables:
"""Test extract_all_tables method for classifying tables."""
def test_extract_all_tables_separates_valid_and_raw_text(self):
"""Test that tables are correctly separated into valid and raw text."""
html = """
| Row 1 text content |
| Row 2 text content |
| Name | Value |
|---|
| Item 1 | 100 |
| Item 2 | 200 |
"""
extractor = HtmlTableExtractor()
valid_tables, raw_text_tables = extractor.extract_all_tables(html)
# First table is single-column (raw text)
# Second table is multi-column (valid data table)
assert len(raw_text_tables) == 1
assert len(valid_tables) == 1
assert raw_text_tables[0].col_count == 1
assert valid_tables[0].col_count == 2
def test_single_column_table_classified_as_raw_text(self):
"""Test that single-column tables are classified as raw text."""
html = """
| This is paragraph text |
| More paragraph text |
| Even more text |
"""
extractor = HtmlTableExtractor()
valid_tables, raw_text_tables = extractor.extract_all_tables(html)
assert len(valid_tables) == 0
assert len(raw_text_tables) == 1
assert raw_text_tables[0].col_count == 1
assert raw_text_tables[0].row_count == 3
def test_multi_column_table_classified_as_valid(self):
"""Test that multi-column tables are classified as valid data tables."""
html = """
"""
extractor = HtmlTableExtractor()
valid_tables, raw_text_tables = extractor.extract_all_tables(html)
assert len(valid_tables) == 1
assert len(raw_text_tables) == 0
assert valid_tables[0].col_count == 3
def test_mixed_tables(self):
"""Test extraction of mixed content: both valid and raw text tables."""
html = """
| Notes section here |
| More notes |
"""
extractor = HtmlTableExtractor()
valid_tables, raw_text_tables = extractor.extract_all_tables(html)
# 2 multi-column tables, 2 single-column tables
assert len(valid_tables) == 2
assert len(raw_text_tables) == 2
def test_empty_html(self):
"""Test extraction from empty HTML."""
extractor = HtmlTableExtractor()
valid_tables, raw_text_tables = extractor.extract_all_tables("")
assert valid_tables == []
assert raw_text_tables == []
def test_no_tables(self):
"""Test extraction from HTML with no tables."""
html = "Just some text
"
extractor = HtmlTableExtractor()
valid_tables, raw_text_tables = extractor.extract_all_tables(html)
assert valid_tables == []
assert raw_text_tables == []
def test_table_index_preserved(self):
"""Test that table indices are preserved during classification."""
html = """
"""
extractor = HtmlTableExtractor()
valid_tables, raw_text_tables = extractor.extract_all_tables(html)
# Valid table was the second table (index 1)
assert valid_tables[0].table_index == 1
# Raw text tables were first and third (indices 0 and 2)
assert raw_text_tables[0].table_index == 0
assert raw_text_tables[1].table_index == 2
class TestExtractTables:
"""Test backward-compatible extract_tables method."""
def test_extract_tables_returns_only_valid(self):
"""Test that extract_tables returns only valid data tables."""
html = """
"""
extractor = HtmlTableExtractor()
tables = extractor.extract_tables(html)
# Should only return the multi-column table
assert len(tables) == 1
assert tables[0].col_count == 2
class TestConvenienceFunctions:
"""Test convenience functions."""
def test_extract_tables_from_html(self):
"""Test extract_tables_from_html convenience function."""
html = """
"""
tables = extract_tables_from_html(html)
# Should only return valid tables
assert len(tables) == 1
assert tables[0].col_count == 2
def test_extract_all_tables_from_html(self):
"""Test extract_all_tables_from_html convenience function."""
html = """
"""
valid_tables, raw_text_tables = extract_all_tables_from_html(html)
assert len(valid_tables) == 1
assert len(raw_text_tables) == 1
class TestTableRowExtraction:
"""Test row extraction from tables."""
def test_raw_text_rows_extracted(self):
"""Test that row content is correctly extracted from raw text tables."""
html = """
| First line of text |
| Second line of text |
| Third line of text |
"""
extractor = HtmlTableExtractor()
_, raw_text_tables = extractor.extract_all_tables(html)
assert len(raw_text_tables) == 1
table = raw_text_tables[0]
assert table.row_count == 3
assert table.rows[0] == ["First line of text"]
assert table.rows[1] == ["Second line of text"]
assert table.rows[2] == ["Third line of text"]
if __name__ == "__main__":
pytest.main([__file__, "-v"])