Image OCR

import subprocess

import demistomock as demisto
import pytest
from CommonServerPython import CommandResults, EntryType
from ImageOCR import (
    CORRUPTED_ERR,
    EMPTY_FILE_MSG,
    extract_text,
    extract_text_command,
    list_languages_command,
    main,
    run_test_module,
)

RETURN_ERROR_TARGET = "ImageOCR.return_error"


def test_list_languages_command():
    """
    When:
     - Running the list_languages function

    Then:
     - Ensure the supported languages in the Docker image are present.
    """
    cmd_res = list_languages_command()
    res = cmd_res.raw_response
    assert len(res) >= 16
    assert "eng" in res  # english
    assert "pol" in res  # polish


@pytest.mark.parametrize(
    "image,expected_text,langs",
    [
        ("irs.png", "Internal Revenue Service", None),
        ("bomb.jpg", "You must transfer bitcoins", None),
        ("noisy1.jpg", "Tesseract OCR", None),
        ("noisy.png", "Tesseract Will", None),
        ("cnbc.gif", "MARKETS", None),
        ("hebrew.tiff", "ביטקוין", ["eng", "heb"]),
    ],
)  # noqa: E124
def test_extract_text(image, expected_text, langs):
    """
    Given:
     - An image with text

    When:
     - Running the extract_text command

    Then:
     - The expected text is extracted
    """
    res = extract_text("test_data/" + image, langs)
    assert expected_text in res


def test_extract_text_verbose_params():
    """
    Given:
     - An image with text

    When:
     - Running the extract_text command

    Then:
     - Validate the result with and without the verbose parameter.
    """
    path = "test_data/bomb.jpg"
    res_verbose = extract_text(path, verbose=True)
    # Some of the verbose data.
    assert "tesseract" in res_verbose
    # Without verbose.
    res_without_verbose = extract_text(path, verbose=False)
    assert "tesseract" not in res_without_verbose


def test_extract_text_command(mocker):
    """
    Given:
     - An image with text

    When:
     - Running the image-ocr-extract-text command

    Then:
     - The expected text is extracted
     - The Human Readable and context are stored with the proper values
    """
    mocker.patch.object(demisto, "args", return_value={"entryid": "test"})
    mocker.patch.object(demisto, "getFilePath", return_value={"path": "test_data/irs.png"})
    mocker.patch.object(demisto, "command", return_value="image-ocr-extract-text")
    mocker.patch.object(demisto, "results")
    # validate our mocks are good
    assert demisto.args()["entryid"] == "test"
    main()
    assert demisto.results.call_count == 1
    # call_args is tuple (args list, kwargs). we only need the first one
    results = demisto.results.call_args[0][0]
    assert results["Type"] == EntryType.NOTE
    assert "Internal Revenue Service" in results["HumanReadable"]
    assert "Internal Revenue Service" in results["EntryContext"]["File(val.EntryID && val.EntryID == obj.EntryID)"]["Text"]


def test_extract_text_command_bad(mocker):
    """
    Given:
     - An image with text
     - a non supported language

    When:
     - Running the image-ocr-extract-text command

    Then:
     - A proper error is raised
    """
    mocker.patch.object(demisto, "args", return_value={"entryid": "test", "langs": "thisis,bad"})
    mocker.patch.object(demisto, "getFilePath", return_value={"path": "test_data/irs.png"})
    mocker.patch.object(demisto, "command", return_value="image-ocr-extract-text")
    return_error_mock = mocker.patch(RETURN_ERROR_TARGET)
    # validate our mocks are good
    assert demisto.args()["entryid"] == "test"
    main()
    assert return_error_mock.call_count == 1
    # call_args last call with a tuple of args list and kwargs
    err_msg = return_error_mock.call_args[0][0]
    assert "Error:" in err_msg
    assert "bad" in err_msg


@pytest.mark.parametrize("skip_corrupted", [True, False])
def test_extract_text_command_corrupted_image(mocker, skip_corrupted: bool):
    """
    Note: if this unittests fails after a docker update, it means tesseract improved corrupted images handling
    Given:
     - A corrupted image
     - The skip_corrupted boolean indicating whether or not to raise an error
    When:
     - Running the image-ocr-extract-text command
    Then:
     - Ensure an error message is returned if skip_corrupted is false, or a warning otherwise.
    """
    mocker.patch.object(demisto, "getFilePath", return_value={"path": "test_data/corrupted.png"})
    results, errors = extract_text_command(
        args={"entryid": "test"},
        instance_languages=["eng"],
        skip_corrupted=skip_corrupted,
    )
    assert len(results + errors) == 1
    if skip_corrupted:
        assert isinstance(results[0], CommandResults)
        assert results[0].entry_type == EntryType.WARNING
    else:
        assert CORRUPTED_ERR in errors[0]


def test_extract_text_command_empty_file(mocker, tmp_path):
    """
    Given:
     - An empty (zero-byte) image file, such as a broken inline email image.

    When:
     - Running the image-ocr-extract-text command.

    Then:
     - Ensure a warning entry is returned instead of an error, and tesseract is never invoked.
    """
    empty_file = tmp_path / "empty.png"
    empty_file.write_bytes(b"")
    mocker.patch.object(demisto, "getFilePath", return_value={"path": str(empty_file)})
    run_mock = mocker.patch("ImageOCR.subprocess.run")

    results, errors = extract_text_command(
        args={"entryid": "test"},
        instance_languages=["eng"],
        skip_corrupted=False,
    )

    assert not errors
    assert len(results) == 1
    assert isinstance(results[0], CommandResults)
    assert results[0].entry_type == EntryType.WARNING
    assert results[0].outputs["Text"] == EMPTY_FILE_MSG
    run_mock.assert_not_called()


@pytest.mark.parametrize("skip_corrupted", [True, False])
def test_extract_text_command_truncated_file(mocker, skip_corrupted: bool):
    """
    Given:
     - A truncated image file that causes tesseract to fail with "truncated file" in stderr.
     - The skip_corrupted boolean indicating whether or not to raise an error.

    When:
     - Running the image-ocr-extract-text command.

    Then:
     - Ensure a warning is returned if skip_corrupted is true, or an error otherwise.
    """
    mocker.patch.object(demisto, "getFilePath", return_value={"path": "test_data/irs.png"})
    mocker.patch("ImageOCR.os.path.getsize", return_value=100)
    mocker.patch(
        "ImageOCR.extract_text",
        side_effect=subprocess.CalledProcessError(
            returncode=1,
            cmd=["tesseract"],
            output="",
            stderr="Error in findFileFormatStream: truncated file",
        ),
    )

    results, errors = extract_text_command(
        args={"entryid": "test"},
        instance_languages=["eng"],
        skip_corrupted=skip_corrupted,
    )

    assert len(results + errors) == 1
    if skip_corrupted:
        assert isinstance(results[0], CommandResults)
        assert results[0].entry_type == EntryType.WARNING
    else:
        assert "truncated file" in errors[0]


def test_run_test_module():
    """
    Given:
     - A param with the supported swedish language

    When:
     - Running the test-module command

    Then:
     - An ok is returned
    """
    res = run_test_module(["swe"])
    assert res == "ok"


def test_run_test_module_bad(mocker):
    """
    Given:
     - A param with the non supported valyrian language

    When:
     - Running the test-module command

    Then:
     - A proper error is presented
    """
    mocker.patch.object(demisto, "params", return_value={"langs": "valyrian"})
    mocker.patch.object(demisto, "command", return_value="test-module")
    mocker.patch.object(demisto, "results")
    return_error_mock = mocker.patch(RETURN_ERROR_TARGET)
    # validate our mocks are good
    assert demisto.command() == "test-module"
    main()
    assert return_error_mock.call_count == 1
    # call_args last call with a tuple of args list and kwargs
    err_msg = return_error_mock.call_args[0][0]
    assert "Unsupported language configured: valyrian" in err_msg