text: # pylint: disable=missing-docstring,invalid-name,protected-access
  import email
  import hashlib
  import logging
  from email import policy
  from pathlib import Path
  from tempfile import TemporaryDirectory
  from unittest.mock import MagicMock, Mock, patch
  
  import pytest
  
  import libratom
  from libratom.data import MIME_TYPES
  from libratom.lib.concurrency import get_messages
  from libratom.lib.constants import SPACY_MODELS, BodyType
  from libratom.lib.core import (
      extract_message_from_archive,
      get_set_of_files,
      load_spacy_model,
      open_mail_archive,
  )
  from libratom.lib.database import db_init, db_session
  from libratom.lib.download import download_files
  from libratom.lib.entities import extract_entities
  from libratom.lib.exceptions import FileTypeError
  from libratom.lib.mbox import MboxArchive
  from libratom.lib.pff import PffArchive
  from libratom.lib.report import generate_report, get_file_info, scan_files
  from libratom.lib.utils import cleanup_message_body
  from libratom.models import FileReport
  
  logger = logging.getLogger(__name__)
  
  
  def test_version():
      assert libratom.__version__
  
  
  def test_pffarchive_load_from_file_object(sample_pst_file):
  
      with sample_pst_file.open(mode="rb") as f, PffArchive(f) as archive:
          assert len(list(archive.messages())) == 2668
  
  
  def test_pffarchive_load_from_invalid_type():
  
      with pytest.raises(TypeError):
          _ = PffArchive(1)
  
  
  def test_open_mail_archive_with_unsupported_type():
  
      with pytest.raises(FileTypeError):
          _ = open_mail_archive(Path("bad_path"))
  
  
  @pytest.mark.parametrize("bfs", [False, True])
  def test_pffarchive_iterate_over_messages(sample_pst_file, bfs):
  
      with PffArchive(sample_pst_file) as archive:
          for message in archive.messages(bfs=bfs):
              assert message.plain_text_body
  
  
  def test_pffarchive_format_message(enron_dataset_part004, empty_message):
  
      for pst_file in enron_dataset_part004.glob("*.pst"):
          with PffArchive(pst_file) as archive:
              for message in archive.messages():
                  # The assertion here doesn't matter as much as
                  # not getting an exception from python's email parsing module
                  assert email.message_from_string(
                      archive.format_message(message), policy=policy.default
                  ) or not archive.format_message(message)
  
      assert PffArchive.format_message(empty_message) == ""
  
  
  def test_get_transport_headers_from_sent_items(enron_dataset_part004):
  
      for pst_file in enron_dataset_part004.glob("*.pst"):
          with PffArchive(pst_file) as archive:
              for folder in archive.folders():
                  try:
                      name = folder.name.lower()
                  except AttributeError:
                      # pylint: disable=no-member
                      if folder.identifier != archive._data.root_folder.identifier:
                          raise
                      continue
                  if "sent mail" in name or "sent items" in name:
                      for message in folder.sub_messages:
                          assert message.transport_headers
  
  
  def test_extract_message_attachments(enron_dataset_part002):
      """Checking 3 known attachments, to validate the attachment extraction process
      """
  
      digests = {
          47685: "d48232614b01e56014293854abbb5db3",
          47717: "cf8be7cd3e6e14307972246e2942c9d1",
          47749: "081e6b66dc89671ff6460adac94dbab1",
      }
  
      with PffArchive(
          next(enron_dataset_part002.glob("*.pst"))
      ) as archive, TemporaryDirectory() as tmp_dir:
  
          # Get message by ID
          node = archive.tree.get_node(2128676)
          message = node.data
  
          for att in message.attachments:
              # Read attachment as bytes
              rbuf = att.read_buffer(att.size)
  
              # Save attachment
              filepath = (
                  Path(tmp_dir) / f"attachment_{message.identifier}_{att.identifier}"
              )
              filepath.write_bytes(rbuf)
  
              # Confirm checksum
              assert hashlib.md5(rbuf).hexdigest() == digests[att.identifier]
  
              # Sanity check on the file
              assert filepath.stat().st_size == att.size
  
  
  def test_get_messages_with_bad_files(enron_dataset_part044, mock_progress_callback):
  
      _count = 0
      for _count, res in enumerate(
          get_messages(
              files=enron_dataset_part044.glob("*.pst"),
              progress_callback=mock_progress_callback,
          ),
          start=1,
      ):
          assert res
  
      assert _count == 558
  
  
  def test_get_message_by_id(sample_pst_file):
      with PffArchive(sample_pst_file) as archive:
          for message in archive.messages():
              msg = archive.get_message_by_id(message.identifier)
              assert msg.identifier == message.identifier
              assert archive.format_message(msg) == archive.format_message(message)
  
  
  def test_get_message_by_id_with_bad_id(sample_pst_file):
      with PffArchive(sample_pst_file) as archive:
          assert archive.get_message_by_id(1234) is None
  
  
  def test_file_report_with_empty_relationship():
      file_report = FileReport()
  
      assert file_report.processing_start_time is None
      assert file_report.processing_end_time is None
      assert file_report.processing_wall_time is None
  
  
  def test_extract_entities_from_mbox_files(directory_of_mbox_files):
  
      tmp_filename = "test.sqlite3"
  
      with TemporaryDirectory() as tmpdir:
  
          destination = Path(tmpdir) / tmp_filename
          Session = db_init(destination)
  
          with db_session(Session) as session:
  
              status = extract_entities(
                  files=get_set_of_files(directory_of_mbox_files),
                  session=session,
                  spacy_model=load_spacy_model(SPACY_MODELS.en_core_web_sm)[0],
                  jobs=2,
              )
  
          assert status == 0
  
  
  @pytest.mark.parametrize(
      "function, patched, kwargs",
      [
          (
              extract_entities,
              "libratom.lib.entities.Message",
              {
                  "spacy_model": load_spacy_model(SPACY_MODELS.en_core_web_sm)[0],
                  "jobs": 2,
              },
          ),
          (generate_report, "libratom.lib.report.Message", {}),
      ],
  )
  def test_run_function_with_interrupt(
      directory_of_mbox_files, function, patched, kwargs
  ):
  
      tmp_filename = "test.sqlite3"
  
      with TemporaryDirectory() as tmpdir:
  
          destination = Path(tmpdir) / tmp_filename
          Session = db_init(destination)
  
          with db_session(Session) as session, patch(
              patched, new=MagicMock(side_effect=KeyboardInterrupt),
          ):
  
              status = function(
                  files=get_set_of_files(directory_of_mbox_files),
                  session=session,
                  **kwargs,
              )
  
          assert status == 1
  
  
  def test_scan_files_with_interrupt(directory_of_mbox_files):
  
      tmp_filename = "test.sqlite3"
  
      with TemporaryDirectory() as tmpdir:
  
          destination = Path(tmpdir) / tmp_filename
          Session = db_init(destination)
  
          with db_session(Session) as session, patch(
              "libratom.lib.report.FileReport",
              new=MagicMock(side_effect=KeyboardInterrupt),
          ):
  
              assert (
                  scan_files(
                      files=get_set_of_files(directory_of_mbox_files),
                      session=session,
                      jobs=2,
                  )
                  == 1
              )
  
  
  @pytest.mark.parametrize("dry_run", [False, True])
  def test_download_files(directory_of_mbox_files, dry_run):
  
      assert directory_of_mbox_files  # so that the files are already present
  
      # Try to re-download files already downloaded by the fixture
      url_template = (
          "https://mail-archives.apache.org/mod_mbox/httpd-users/20190{month}.mbox"
      )
      path = Path("/tmp/libratom/test_data/httpd-users")
      urls = [url_template.format(month=i) for i in range(1, 7)]
      download_files(urls, path, dry_run=dry_run)
  
  
  def test_download_files_with_bad_urls():
  
      bad_urls = ["http://foobar"] * 6
  
      with TemporaryDirectory() as tmpdir, patch("requests.Session.get") as mock_get:
          mock_get.return_value.ok = False
  
          with pytest.raises(RuntimeError):
              download_files(bad_urls, Path(tmpdir))
  
  
  def test_utf8_message_with_no_cte_header_as_string(utf8_message_with_no_cte_header):
      assert MboxArchive.format_message(utf8_message_with_no_cte_header)
  
  
  def test_get_file_info(sample_pst_file):
  
      res, error = get_file_info(
          # Must use dictionary form if function is called explicitly
          {"path": sample_pst_file}
      )
  
      assert res.get("path") == str(sample_pst_file)
      assert res.get("name") == sample_pst_file.name
      assert res.get("size") == 172450816
      assert res.get("md5") == "1038b99c2c323ca563da79dbbee3876f"
      assert (
          res.get("sha256")
          == "12c10d46f3935680f5c849a66737cf442fc171a2a20dce32c5efcc668366be96"
      )
      assert res.get("msg_count") == 2668
      assert not error
  
  
  def test_attachments_mime_type_validation(enron_dataset, mock_progress_callback):
  
      files = get_set_of_files(enron_dataset)
  
      for res in get_messages(files, progress_callback=mock_progress_callback):
          attachments = res.get("attachments")
          if attachments:
              for attachment in attachments:
                  try:
                      assert attachment.mime_type in MIME_TYPES
                  except AssertionError:
                      # Some enron files have these obsolete attachment types
                      assert attachment.mime_type in [
                          "application/msexcell",
                          "application/mspowerpoint",
                      ]
  
  
  def test_get_mbox_message_by_id(sample_mbox_file):
      with open_mail_archive(sample_mbox_file) as archive:
  
          assert archive.message_count == 113
  
          for index, message in enumerate(archive.messages(), start=1):
              msg = archive.get_message_by_id(index)
              assert extract_message_from_archive(archive, index)
              assert archive.format_message(msg) == archive.format_message(message)
              assert archive.get_message_headers(message)
  
  
  def test_get_mbox_message_by_id_with_bad_id(sample_mbox_file):
      with open_mail_archive(sample_mbox_file) as archive:
          assert archive.get_message_by_id(1234) is None
  
  
  @pytest.mark.parametrize("mock_cls", [MagicMock, Mock])
  def test_get_attachment_metadata(mock_cls):
      message = MagicMock(identifier=123, attachments=[mock_cls(name="foo", size="0")])
  
      assert PffArchive().get_attachment_metadata(message)[0].mime_type is None
  
  
  @pytest.mark.parametrize(
      "message, body_type",
      [
          (Mock(), BodyType.PLAIN),
          (Mock(plain_text_body=None), BodyType.RTF),
          (Mock(plain_text_body=None, rtf_body=None), BodyType.HTML),
          (Mock(plain_text_body=None, rtf_body=None, html_body=None), None),
      ],
  )
  def test_get_message_body(message, body_type):
      assert PffArchive().get_message_body(message)[1] is body_type
  
  
  @pytest.mark.parametrize(
      "body, body_type, result",
      [
          (
              r"{\rtf {\fonttbl {\f0 Times New Roman;}} \f0\fs60 foo} ",
              BodyType.RTF,
              "foo",
          ),
          ("<body><table><tr><td>foo</td></tr></table></body>", BodyType.HTML, "foo"),
      ],
  )
  def test_cleanup_message_body(body, body_type, result):
      assert cleanup_message_body(body, body_type) == result
  