test_extensions.py

import tempfile
import os
import struct
import zlib

import pytest
from openpyxl import Workbook, load_workbook

from officecli.handlers.excel_extensions import (
    ConditionalFormatHelper, ValidationHelper, ImageHelper,
    PrintSettingsHelper, ProtectionHelper,
)
from officecli.core.exceptions import CliException, InvalidValueException


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture(scope="function")
def workbook_data():
    """Create a shared workbook for all tests."""
    wb = Workbook()
    ws = wb.active
    ws.title = "T"
    for r in range(1, 11):
        ws.cell(row=r, column=1, value=r * 10)
        ws.cell(row=r, column=2, value=["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"][(r - 1) % 10])

    tmp = tempfile.mktemp(suffix=".xlsx")
    wb.save(tmp)
    wb2 = load_workbook(tmp)
    yield tmp, wb2

    # Cleanup
    if os.path.exists(tmp):
        os.unlink(tmp)


def _make_png(path: str) -> None:
    """Create a minimal 10x10 red PNG for testing."""
    hdr = b"\x89PNG\r\n\x1a\n"

    def chunk(ct, d):
        c = ct + d
        return struct.pack(">I", len(d)) + c + struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)

    ihdr = struct.pack(">IIBBBBB", 10, 10, 8, 2, 0, 0, 0)
    raw = b"".join(b"\x00" + bytes([255, 0, 0]) * 10 for _ in range(10))
    data = hdr + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b"")
    with open(path, "wb") as f:
        f.write(data)


# ---------------------------------------------------------------------------
# Conditional Format Tests
# ---------------------------------------------------------------------------

class TestConditionalFormat:
    """Tests for ConditionalFormatHelper."""

    def test_add_and_count(self, workbook_data):
        tmp, wb = workbook_data
        cfh = ConditionalFormatHelper()
        cfh.add_cell_is_rule(wb, "/sheet[1]", "A1:A10", "greaterThan", "50", fill_color="FFCCCC")
        cfh.add_color_scale(wb, "/sheet[1]", "A1:A10", "63BE7B", "F8696B", mid_color="FFEB84")
        cfh.add_data_bar(wb, "/sheet[1]", "A1:A10", color="638EC6")
        cfh.add_icon_set(wb, "/sheet[1]", "A1:A10", icon_style="3TrafficLights1")
        cfh.add_icon_set(wb, "/sheet[1]", "B1:B10", icon_style="5Arrows")
        cfh.add_top10_rule(wb, "/sheet[1]", "A1:A10", rank=3, fill_color="FFFFCC")
        cfh.add_above_average(wb, "/sheet[1]", "A1:A10", fill_color="C6EFCE")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "containsText", "A", fill_color="FFC000")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "notContainsText", "Z", fill_color="E2EFDA")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "beginsWith", "A", fill_color="FFC000")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "endsWith", "B", fill_color="FFC000")
        cfh.add_duplicate_rule(wb, "/sheet[1]", "B1:B10", duplicate=True, fill_color="D9E1F2")
        cfh.add_duplicate_rule(wb, "/sheet[1]", "B1:B10", duplicate=False, fill_color="D9E1F2")
        cfh.add_formula_rule(wb, "/sheet[1]", "A1:A10", "A1>55", fill_color="B4C6E7")
        cfh.add_time_period_rule(wb, "/sheet[1]", "A1:A10", "today", fill_color="E2EFDA")
        cfh.add_blank_rule(wb, "/sheet[1]", "A1:A10", blanks=True, fill_color="F2F2F2")
        cfh.add_blank_rule(wb, "/sheet[1]", "A1:A10", blanks=False, fill_color="E2EFDA")
        cfh.add_cell_is_rule(wb, "/sheet[1]", "A1:A10", "between", ["20", "60"], fill_color="BDD7EE")
        n = len(cfh.list_conditional_formats(wb, "/sheet[1]"))
        assert n == 18

    def test_persist_and_reload(self, workbook_data):
        tmp, wb = workbook_data
        cfh = ConditionalFormatHelper()
        # Re-add rules (independent from other tests)
        cfh.add_cell_is_rule(wb, "/sheet[1]", "A1:A10", "greaterThan", "50", fill_color="FFCCCC")
        cfh.add_color_scale(wb, "/sheet[1]", "A1:A10", "63BE7B", "F8696B", mid_color="FFEB84")
        cfh.add_data_bar(wb, "/sheet[1]", "A1:A10", color="638EC6")
        cfh.add_icon_set(wb, "/sheet[1]", "A1:A10", icon_style="3TrafficLights1")
        cfh.add_icon_set(wb, "/sheet[1]", "B1:B10", icon_style="5Arrows")
        cfh.add_top10_rule(wb, "/sheet[1]", "A1:A10", rank=3, fill_color="FFFFCC")
        cfh.add_above_average(wb, "/sheet[1]", "A1:A10", fill_color="C6EFCE")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "containsText", "A", fill_color="FFC000")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "notContainsText", "Z", fill_color="E2EFDA")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "beginsWith", "A", fill_color="FFC000")
        cfh.add_text_rule(wb, "/sheet[1]", "B1:B10", "endsWith", "B", fill_color="FFC000")
        cfh.add_duplicate_rule(wb, "/sheet[1]", "B1:B10", duplicate=True, fill_color="D9E1F2")
        cfh.add_duplicate_rule(wb, "/sheet[1]", "B1:B10", duplicate=False, fill_color="D9E1F2")
        cfh.add_formula_rule(wb, "/sheet[1]", "A1:A10", "A1>55", fill_color="B4C6E7")
        cfh.add_time_period_rule(wb, "/sheet[1]", "A1:A10", "today", fill_color="E2EFDA")
        cfh.add_blank_rule(wb, "/sheet[1]", "A1:A10", blanks=True, fill_color="F2F2F2")
        cfh.add_blank_rule(wb, "/sheet[1]", "A1:A10", blanks=False, fill_color="E2EFDA")
        cfh.add_cell_is_rule(wb, "/sheet[1]", "A1:A10", "between", ["20", "60"], fill_color="BDD7EE")

        wb.save(tmp)
        wb3 = load_workbook(tmp)
        n2 = len(cfh.list_conditional_formats(wb3, "/sheet[1]"))
        assert n2 == 18

    def test_remove_and_clear(self, workbook_data):
        tmp, wb = workbook_data
        cfh = ConditionalFormatHelper()
        cfh.add_cell_is_rule(wb, "/sheet[1]", "A1:A10", "greaterThan", "50", fill_color="FFCCCC")
        cfh.add_color_scale(wb, "/sheet[1]", "A1:A10", "63BE7B", "F8696B")

        assert len(cfh.list_conditional_formats(wb, "/sheet[1]")) >= 2
        cfh.remove_conditional_format(wb, "/sheet[1]", 0)
        cfh.clear_conditional_formats(wb, "/sheet[1]")
        assert len(cfh.list_conditional_formats(wb, "/sheet[1]")) == 0

    def test_invalid_sheet_path(self, workbook_data):
        tmp, wb = workbook_data
        cfh = ConditionalFormatHelper()
        with pytest.raises(CliException):
            cfh.add_cell_is_rule(wb, "/sheet[999]", "A1", "equal", "1")

    def test_invalid_operator(self, workbook_data):
        tmp, wb = workbook_data
        cfh = ConditionalFormatHelper()
        with pytest.raises(InvalidValueException):
            cfh.add_cell_is_rule(wb, "/sheet[1]", "A1", "badOp", "1")

    def test_invalid_color(self, workbook_data):
        tmp, wb = workbook_data
        cfh = ConditionalFormatHelper()
        with pytest.raises(InvalidValueException):
            cfh.add_color_scale(wb, "/sheet[1]", "A1", "badcolor", "FF0000")


# ---------------------------------------------------------------------------
# Validation Tests
# ---------------------------------------------------------------------------

class TestValidation:
    """Tests for ValidationHelper."""

    def test_add_and_list(self, workbook_data):
        tmp, wb = workbook_data
        vh = ValidationHelper()
        vh.add_validation(wb, "/sheet[1]", "C1:C10", validation_type="list", formula1="Yes,No,Maybe")
        vh.add_validation(wb, "/sheet[1]", "D1:D10", validation_type="whole", operator="greaterThan", formula1="0")
        assert len(vh.list_validations(wb, "/sheet[1]")) == 2

    def test_remove_and_clear(self, workbook_data):
        tmp, wb = workbook_data
        vh = ValidationHelper()
        vh.add_validation(wb, "/sheet[1]", "C1:C10", validation_type="list", formula1="Yes,No,Maybe")
        vh.add_validation(wb, "/sheet[1]", "D1:D10", validation_type="whole", operator="greaterThan", formula1="0")
        vh.remove_validation(wb, "/sheet[1]", 0)
        vh.clear_validations(wb, "/sheet[1]")
        assert len(vh.list_validations(wb, "/sheet[1]")) == 0

    def test_invalid_type(self, workbook_data):
        tmp, wb = workbook_data
        vh = ValidationHelper()
        with pytest.raises(InvalidValueException):
            vh.add_validation(wb, "/sheet[1]", "A1", validation_type="badType")


# ---------------------------------------------------------------------------
# Print Settings Tests
# ---------------------------------------------------------------------------

class TestPrintSettings:
    """Tests for PrintSettingsHelper."""

    def test_page_setup_and_margins(self, workbook_data):
        tmp, wb = workbook_data
        psh = PrintSettingsHelper()
        psh.set_page_setup(wb, "/sheet[1]", orientation="landscape", paper_size="A4")
        psh.set_margins(wb, "/sheet[1]", left=0.5, right=0.5)
        s = psh.get_page_setup(wb, "/sheet[1]")
        assert s["orientation"] == "landscape"
        assert s["paper_size"] == "a4"

    def test_print_area_and_titles(self, workbook_data):
        tmp, wb = workbook_data
        psh = PrintSettingsHelper()
        psh.set_print_area(wb, "/sheet[1]", "A1:C10")
        psh.set_print_title_rows(wb, "/sheet[1]", "$1:$1")
        psh.set_print_title_cols(wb, "/sheet[1]", "$A:$A")
        psh.set_page_breaks(wb, "/sheet[1]", row_breaks=[5], col_breaks=[2])
        # Verify print area is set
        s = psh.get_page_setup(wb, "/sheet[1]")
        assert s["print_area"]

    def test_clear_print_area(self, workbook_data):
        tmp, wb = workbook_data
        psh = PrintSettingsHelper()
        psh.set_print_area(wb, "/sheet[1]", "A1:C10")
        psh.clear_print_area(wb, "/sheet[1]")
        assert not psh.get_page_setup(wb, "/sheet[1]")["print_area"]

    def test_invalid_paper_size(self, workbook_data):
        tmp, wb = workbook_data
        psh = PrintSettingsHelper()
        with pytest.raises(InvalidValueException):
            psh.set_page_setup(wb, "/sheet[1]", paper_size="BadSize")

    def test_invalid_margin(self, workbook_data):
        tmp, wb = workbook_data
        psh = PrintSettingsHelper()
        with pytest.raises(InvalidValueException):
            psh.set_margins(wb, "/sheet[1]", left=-1)


# ---------------------------------------------------------------------------
# Protection Tests
# ---------------------------------------------------------------------------

class TestProtection:
    """Tests for ProtectionHelper."""

    def test_protect_and_unprotect_sheet(self, workbook_data):
        tmp, wb = workbook_data
        ph = ProtectionHelper()
        ph.protect_sheet(wb, "/sheet[1]", password="x", allow_sort=True, allow_autofilter=True)
        assert ph.is_sheet_protected(wb, "/sheet[1]")
        info = ph.get_sheet_protection_info(wb, "/sheet[1]")
        assert info["protected"] and info["allow_sort"] and info["allow_autofilter"]
        ph.unprotect_sheet(wb, "/sheet[1]")
        assert not ph.is_sheet_protected(wb, "/sheet[1]")

    def test_protect_and_unprotect_workbook(self, workbook_data):
        tmp, wb = workbook_data
        ph = ProtectionHelper()
        ph.protect_workbook(wb, password="y", lock_structure=True)
        assert ph.is_workbook_protected(wb)
        ph.unprotect_workbook(wb)
        assert not ph.is_workbook_protected(wb)


# ---------------------------------------------------------------------------
# Image Tests
# ---------------------------------------------------------------------------

class TestImage:
    """Tests for ImageHelper."""

    @pytest.fixture
    def png_file(self, workbook_data):
        tmp, wb = workbook_data
        img = tmp + ".png"
        _make_png(img)
        yield tmp, wb, img
        if os.path.exists(img):
            os.unlink(img)

    def test_add_resize_remove(self, png_file):
        tmp, wb, img = png_file
        ih = ImageHelper()
        ih.add_image(wb, "/sheet[1]", img, anchor="C1", width=100, height=100)
        assert len(ih.list_images(wb, "/sheet[1]")) == 1
        ih.resize_image(wb, "/sheet[1]", 0, width=200, height=50)
        ih.remove_image(wb, "/sheet[1]", 0)
        assert len(ih.list_images(wb, "/sheet[1]")) == 0

    def test_nonexistent_file(self, png_file):
        tmp, wb, img = png_file
        ih = ImageHelper()
        with pytest.raises(CliException):
            ih.add_image(wb, "/sheet[1]", "/nonexistent.png")