Skip to content

Commit

Permalink
Attach uploaded documents to pno emails
Browse files Browse the repository at this point in the history
  • Loading branch information
@VincentAntoine
VincentAntoine committed Sep 25, 2024
1 parent 9824b80 commit 5fbbb33
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 51 deletions.
33 changes: 20 additions & 13 deletions datascience/src/pipeline/flows/distribute_pnos.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from jinja2 import Environment, FileSystemLoader, Template, select_autoescape
from prefect import Flow, Parameter, case, flatten, task, unmapped
from prefect.executors import LocalDaskExecutor
from sqlalchemy import Executable, bindparam, select, text
from sqlalchemy import Executable, Select, bindparam, select, text

from config import (
CNSP_CROSSA_CACEM_LOGOS_PATH,
Expand Down Expand Up @@ -126,7 +126,7 @@ def extract_facade_email_addresses() -> dict:
@task(checkpoint=False)
def make_prior_notification_attachments_query(
pnos: List[RenderedPno], prior_notification_uploads_table: sqlalchemy.Table
) -> sqlalchemy.Select | None:
) -> Select | None:
pno_ids = sorted(set([pno.report_id for pno in pnos]))

if pno_ids:
Expand All @@ -143,12 +143,12 @@ def make_prior_notification_attachments_query(

@task(checkpoint=False)
def execute_prior_notification_attachments_query(
query: sqlalchemy.Executable | None,
query: Select | None,
) -> dict:
if isinstance(query, sqlalchemy.Executable):
if isinstance(query, Select):
attachments = read_query(query, db="monitorfish_remote")
attachments["filename_content"] = attachments.apply(
lambda row: (row["file_name"], row["content"].tobytes()), axis=1
lambda row: (row["file_name"], row["content"]), axis=1
)
res = attachments.groupby("report_id")["filename_content"].agg(list).to_dict()
else:
Expand Down Expand Up @@ -698,7 +698,9 @@ def attribute_addressees(


@task(checkpoint=False)
def create_email(pno: RenderedPno, test_mode: bool) -> PnoToSend:
def create_email(
pno: RenderedPno, uploaded_attachments: dict, test_mode: bool
) -> PnoToSend:
if pno.emails:
if test_mode:
if PNO_TEST_EMAIL:
Expand All @@ -708,6 +710,13 @@ def create_email(pno: RenderedPno, test_mode: bool) -> PnoToSend:
else:
to = pno.emails

attachments = [
(
f"Preavis_{pno.vessel_name if pno.vessel_name else ''}.pdf",
pno.pdf_document,
)
] + (uploaded_attachments.get(pno.report_id) or [])

message = create_html_email(
to=to,
subject=f"Préavis de débarquement - {pno.vessel_name}",
Expand All @@ -718,12 +727,7 @@ def create_email(pno: RenderedPno, test_mode: bool) -> PnoToSend:
LIBERTE_EGALITE_FRATERNITE_LOGO_PATH,
MARIANNE_LOGO_PATH,
],
attachments=[
(
f"Preavis_{pno.vessel_name if pno.vessel_name else ''}.pdf",
pno.pdf_document,
)
],
attachments=attachments,
reply_to=CNSP_FRANCE_EMAIL_ADDRESS,
)
return PnoToSend(
Expand Down Expand Up @@ -1072,7 +1076,9 @@ def make_manual_prior_notifications_statement(
)

email = create_email.map(
pnos_with_addressees, test_mode=unmapped(test_mode)
pnos_with_addressees,
uploaded_attachments=unmapped(attachments),
test_mode=unmapped(test_mode),
)
email = filter_results(email)

Expand Down Expand Up @@ -1120,6 +1126,7 @@ def make_manual_prior_notifications_statement(

flow.set_reference_tasks(
[
attachments,
pnos_to_generate,
pnos_to_distribute,
pnos_with_addressees,
Expand Down
101 changes: 63 additions & 38 deletions datascience/tests/test_pipeline/test_flows/test_distribute_pnos.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,16 @@
from dateutil.relativedelta import relativedelta
from jinja2 import Template
from requests import Response
from sqlalchemy import BOOLEAN, TIMESTAMP, VARCHAR, Column, MetaData, Table
from sqlalchemy import (
BOOLEAN,
TIMESTAMP,
VARCHAR,
Column,
MetaData,
Select,
Table,
select,
)
from sqlalchemy.dialects.postgresql import BYTEA

from config import EMAIL_IMAGES_LOCATION, TEST_DATA_LOCATION, default_risk_factors
Expand Down Expand Up @@ -1276,31 +1285,25 @@ def prior_notification_uploads_table() -> Table:


@pytest.fixture
def pnos_query_string() -> str:
return (
"SELECT "
"public.prior_notification_uploads.report_id, "
"public.prior_notification_uploads.file_name, "
"public.prior_notification_uploads.content "
"\nFROM public.prior_notification_uploads "
"\nWHERE "
"public.prior_notification_uploads.report_id IN "
"('00000000-0000-4000-0000-000000000003', '11', '456-def', '789-ghi')"
)
def pnos_query(prior_notification_uploads_table) -> Select:
pno_ids = ["00000000-0000-4000-0000-000000000003", "11", "456-def", "789-ghi"]

return select(
prior_notification_uploads_table.c.report_id,
prior_notification_uploads_table.c.file_name,
prior_notification_uploads_table.c.content,
).where(prior_notification_uploads_table.c.report_id.in_(pno_ids))


@pytest.fixture
def pnos_query_string_empty_result() -> str:
return (
"SELECT "
"public.prior_notification_uploads.report_id, "
"public.prior_notification_uploads.file_name, "
"public.prior_notification_uploads.content "
"\nFROM public.prior_notification_uploads "
"\nWHERE "
"public.prior_notification_uploads.report_id IN "
"('Does', 'Not', 'Exist')"
)
def pnos_query_empty_result(prior_notification_uploads_table) -> Select:
pno_ids = ["Does", "Not", "Exist"]

return select(
prior_notification_uploads_table.c.report_id,
prior_notification_uploads_table.c.file_name,
prior_notification_uploads_table.c.content,
).where(prior_notification_uploads_table.c.report_id.in_(pno_ids))


@pytest.fixture
Expand Down Expand Up @@ -1722,14 +1725,15 @@ def test_attribute_addressees_when_is_verified(


def test_make_prior_notification_attachments_query(
rendered_pnos, prior_notification_uploads_table, pnos_query_string
rendered_pnos, prior_notification_uploads_table, pnos_query
):
# With non empty input
query = make_prior_notification_attachments_query.run(
rendered_pnos, prior_notification_uploads_table
)
query_string = str(query.compile(compile_kwargs={"literal_binds": True}))
assert query_string == pnos_query_string
assert str(query.compile(compile_kwargs={"literal_binds": True})) == str(
pnos_query.compile(compile_kwargs={"literal_binds": True})
)

# With empty input
query = make_prior_notification_attachments_query.run(
Expand All @@ -1739,33 +1743,39 @@ def test_make_prior_notification_attachments_query(


def test_execute_prior_notification_attachments_query(
reset_test_data, pnos_query_string, prior_notification_attachments
reset_test_data, pnos_query, prior_notification_attachments
):
attachments = execute_prior_notification_attachments_query.run(
sqlalchemy.text(pnos_query_string)
)
attachments = execute_prior_notification_attachments_query.run(pnos_query)
assert attachments == prior_notification_attachments


def test_execute_prior_notification_attachments_query_when_result_is_empty(
reset_test_data, pnos_query_string_empty_result
reset_test_data, pnos_query_empty_result
):
attachments = execute_prior_notification_attachments_query.run(
sqlalchemy.text(pnos_query_string_empty_result)
pnos_query_empty_result
)
assert attachments == dict()


@pytest.mark.parametrize("test_mode", [False, True])
@pytest.mark.parametrize(
"test_mode,pno_has_uploaded_attachments",
[(False, False), (False, True), (True, False), (True, True)],
)
def test_create_email(
pno_pdf_document_to_distribute_targeted_vessel_and_segments_assigned,
cnsp_crossa_cacem_logos,
marianne_gif,
liberte_egalite_fraternite_gif,
prior_notification_attachments,
test_mode,
pno_has_uploaded_attachments,
):
pno_to_send = create_email.run(
pno_pdf_document_to_distribute_targeted_vessel_and_segments_assigned,
uploaded_attachments=prior_notification_attachments
if pno_has_uploaded_attachments
else dict(),
test_mode=test_mode,
)

Expand All @@ -1792,10 +1802,22 @@ def test_create_email(
)

attachments = list(pno_to_send.message.iter_attachments())
assert len(attachments) == 1
attachment = attachments[0]
assert attachment.get_content_type() == "application/octet-stream"
assert attachment.get_content() == b"PDF Document"
assert len(attachments) == 3 if pno_has_uploaded_attachments else 1
attachment_0 = attachments[0]
assert attachment_0.get_filename() == "Preavis_Le navire 123-abc.pdf"
assert attachment_0.get_content_type() == "application/octet-stream"
assert attachment_0.get_content() == b"PDF Document"

if pno_has_uploaded_attachments:
attachment_1 = attachments[1]
assert attachment_1.get_filename() == "Uploaded document n°1.pdf"
assert attachment_1.get_content_type() == "application/octet-stream"
assert attachment_1.get_content() == b"PDF"

attachment_2 = attachments[2]
assert attachment_2.get_filename() == "Uploaded document n°2.docx"
assert attachment_2.get_content_type() == "application/octet-stream"
assert attachment_2.get_content() == b"Text Document"

body = pno_to_send.message.get_body()
assert body.get_content_type() == "multipart/related"
Expand Down Expand Up @@ -1828,14 +1850,17 @@ def test_create_email(

def test_create_email_with_no_email_addressees_returns_none(
pno_pdf_document_to_distribute_without_addressees_assigned,
prior_notification_attachments,
):
pno_to_send = create_email.run(
pno_pdf_document_to_distribute_without_addressees_assigned, test_mode=False
)
assert pno_to_send is None

pno_to_send = create_email.run(
pno_pdf_document_to_distribute_without_addressees_assigned, test_mode=True
pno_pdf_document_to_distribute_without_addressees_assigned,
uploaded_attachments=prior_notification_attachments,
test_mode=True,
)
assert pno_to_send is None

Expand Down

0 comments on commit 5fbbb33

Please sign in to comment.