From 7dcf18ea1e5bfcab355f3dd8c356a01861dcadbd Mon Sep 17 00:00:00 2001 From: Max Chis Date: Fri, 26 Sep 2025 19:44:34 -0400 Subject: [PATCH] Add annotation logic for new agency suggestion --- ...3e23d3f0_update_locations_expanded_view.py | 1 + .../annotate/all/post/models/agency.py | 18 ++++++ .../annotate/all/post/models/request.py | 11 ++-- src/api/endpoints/annotate/all/post/query.py | 4 +- .../endpoints/annotate/all/post/requester.py | 16 +++++ .../models/impl/agency/suggestion/__init__.py | 0 .../models/impl/agency/suggestion/pydantic.py | 17 +++++ .../impl/agency/suggestion/sqlalchemy.py | 19 ++++++ .../api/annotate/all/test_happy_path.py | 5 +- .../api/annotate/all/test_new_agency.py | 64 +++++++++++++++++++ .../annotate/all/test_post_batch_filtering.py | 3 +- .../api/annotate/all/test_validation_error.py | 3 +- .../unit/api/test_all_annotation_post_info.py | 5 +- 13 files changed, 155 insertions(+), 11 deletions(-) create mode 100644 src/api/endpoints/annotate/all/post/models/agency.py create mode 100644 src/db/models/impl/agency/suggestion/__init__.py create mode 100644 src/db/models/impl/agency/suggestion/pydantic.py create mode 100644 src/db/models/impl/agency/suggestion/sqlalchemy.py create mode 100644 tests/automated/integration/api/annotate/all/test_new_agency.py diff --git a/alembic/versions/2025_09_26_1751-d4c63e23d3f0_update_locations_expanded_view.py b/alembic/versions/2025_09_26_1751-d4c63e23d3f0_update_locations_expanded_view.py index 675fd7b2..871e54b9 100644 --- a/alembic/versions/2025_09_26_1751-d4c63e23d3f0_update_locations_expanded_view.py +++ b/alembic/versions/2025_09_26_1751-d4c63e23d3f0_update_locations_expanded_view.py @@ -38,6 +38,7 @@ def _update_locations_expanded_view(): WHEN locations.type = 'Locality'::location_type THEN localities.name WHEN locations.type = 'County'::location_type THEN counties.name::character varying WHEN locations.type = 'State'::location_type THEN us_states.state_name::character varying + WHEN locations.type = 'National'::location_type THEN 'United States' ELSE NULL::character varying END AS display_name, CASE diff --git a/src/api/endpoints/annotate/all/post/models/agency.py b/src/api/endpoints/annotate/all/post/models/agency.py new file mode 100644 index 00000000..55c52e49 --- /dev/null +++ b/src/api/endpoints/annotate/all/post/models/agency.py @@ -0,0 +1,18 @@ +from pydantic import BaseModel + +from src.db.models.impl.agency.enums import JurisdictionType, AgencyType + + +class AnnotationNewAgencySuggestionInfo(BaseModel): + name: str + location_id: int + jurisdiction_type: JurisdictionType | None + agency_type: AgencyType | None + +class AnnotationPostAgencyInfo(BaseModel): + new_agency_suggestion: AnnotationNewAgencySuggestionInfo | None = None + agency_ids: list[int] = [] + + @property + def empty(self) -> bool: + return self.new_agency_suggestion is None and len(self.agency_ids) == 0 diff --git a/src/api/endpoints/annotate/all/post/models/request.py b/src/api/endpoints/annotate/all/post/models/request.py index 939eafab..240c8389 100644 --- a/src/api/endpoints/annotate/all/post/models/request.py +++ b/src/api/endpoints/annotate/all/post/models/request.py @@ -1,5 +1,6 @@ -from pydantic import BaseModel, model_validator +from pydantic import BaseModel, model_validator, ConfigDict +from src.api.endpoints.annotate.all.post.models.agency import AnnotationPostAgencyInfo from src.api.endpoints.annotate.all.post.models.name import AnnotationPostNameInfo from src.core.enums import RecordType from src.core.exceptions import FailedValidationException @@ -7,9 +8,11 @@ class AllAnnotationPostInfo(BaseModel): + model_config = ConfigDict(extra='forbid') + suggested_status: URLType record_type: RecordType | None = None - agency_ids: list[int] + agency_info: AnnotationPostAgencyInfo location_ids: list[int] name_info: AnnotationPostNameInfo = AnnotationPostNameInfo() @@ -30,8 +33,8 @@ def forbid_all_else_if_not_relevant(self): return self if self.record_type is not None: raise FailedValidationException("record_type must be None if suggested_status is NOT RELEVANT") - if len(self.agency_ids) > 0: - raise FailedValidationException("agency_ids must be empty if suggested_status is NOT RELEVANT") + if not self.agency_info.empty: + raise FailedValidationException("agency_info must be empty if suggested_status is NOT RELEVANT") if len(self.location_ids) > 0: raise FailedValidationException("location_ids must be empty if suggested_status is NOT RELEVANT") return self diff --git a/src/api/endpoints/annotate/all/post/query.py b/src/api/endpoints/annotate/all/post/query.py index a27d6c6f..01c6973e 100644 --- a/src/api/endpoints/annotate/all/post/query.py +++ b/src/api/endpoints/annotate/all/post/query.py @@ -46,4 +46,6 @@ async def run(self, session: AsyncSession) -> None: # TODO (TEST): Add test for submitting Meta URL validation requester.optionally_add_record_type(self.post_info.record_type) - requester.add_agency_ids(self.post_info.agency_ids) + requester.add_agency_ids(self.post_info.agency_info.agency_ids) + + await requester.optionally_add_new_agency_suggestion(self.post_info.agency_info.new_agency_suggestion) diff --git a/src/api/endpoints/annotate/all/post/requester.py b/src/api/endpoints/annotate/all/post/requester.py index 44f0e0f7..9f9d0a78 100644 --- a/src/api/endpoints/annotate/all/post/requester.py +++ b/src/api/endpoints/annotate/all/post/requester.py @@ -1,7 +1,9 @@ from sqlalchemy.ext.asyncio import AsyncSession +from src.api.endpoints.annotate.all.post.models.agency import AnnotationNewAgencySuggestionInfo from src.api.endpoints.annotate.all.post.models.name import AnnotationPostNameInfo from src.core.enums import RecordType +from src.db.models.impl.agency.suggestion.sqlalchemy import NewAgencySuggestion from src.db.models.impl.flag.url_validated.enums import URLType from src.db.models.impl.link.user_name_suggestion.sqlalchemy import LinkUserNameSuggestion from src.db.models.impl.url.suggestion.agency.user import UserUrlAgencySuggestion @@ -93,3 +95,17 @@ async def optionally_add_name_suggestion( suggestion_id=name_suggestion.id, ) self.session.add(link) + + async def optionally_add_new_agency_suggestion( + self, + suggestion_info: AnnotationNewAgencySuggestionInfo | None + ) -> None: + if suggestion_info is None: + return + new_agency_suggestion = NewAgencySuggestion( + name=suggestion_info.name, + location_id=suggestion_info.location_id, + jurisdiction_type=suggestion_info.jurisdiction_type, + agency_type=suggestion_info.agency_type, + ) + self.session.add(new_agency_suggestion) \ No newline at end of file diff --git a/src/db/models/impl/agency/suggestion/__init__.py b/src/db/models/impl/agency/suggestion/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/db/models/impl/agency/suggestion/pydantic.py b/src/db/models/impl/agency/suggestion/pydantic.py new file mode 100644 index 00000000..84046717 --- /dev/null +++ b/src/db/models/impl/agency/suggestion/pydantic.py @@ -0,0 +1,17 @@ +from pydantic import BaseModel + +from src.db.models.impl.agency.enums import JurisdictionType, AgencyType +from src.db.models.impl.agency.suggestion.sqlalchemy import NewAgencySuggestion +from src.db.models.templates_.base import Base + + +class NewAgencySuggestionPydantic(BaseModel): + + name: str + location_id: int + jurisdiction_type: JurisdictionType | None + agency_type: AgencyType | None + + @classmethod + def sa_model(cls) -> type[Base]: + return NewAgencySuggestion \ No newline at end of file diff --git a/src/db/models/impl/agency/suggestion/sqlalchemy.py b/src/db/models/impl/agency/suggestion/sqlalchemy.py new file mode 100644 index 00000000..f15b2ef0 --- /dev/null +++ b/src/db/models/impl/agency/suggestion/sqlalchemy.py @@ -0,0 +1,19 @@ +from sqlalchemy import String, Column + +from src.db.models.helpers import enum_column +from src.db.models.impl.agency.enums import JurisdictionType, AgencyType +from src.db.models.mixins import CreatedAtMixin, LocationDependentMixin +from src.db.models.templates_.with_id import WithIDBase + + +class NewAgencySuggestion( + WithIDBase, + CreatedAtMixin, + LocationDependentMixin, +): + + __tablename__ = 'new_agency_suggestions' + + name = Column(String) + jurisdiction_type = enum_column(JurisdictionType, name='jurisdiction_type_enum', nullable=True) + agency_type = enum_column(AgencyType, name='agency_type_enum', nullable=True) \ No newline at end of file diff --git a/tests/automated/integration/api/annotate/all/test_happy_path.py b/tests/automated/integration/api/annotate/all/test_happy_path.py index c7e1c5b5..7721e80c 100644 --- a/tests/automated/integration/api/annotate/all/test_happy_path.py +++ b/tests/automated/integration/api/annotate/all/test_happy_path.py @@ -3,6 +3,7 @@ from src.api.endpoints.annotate.all.get.models.location import LocationAnnotationUserSuggestion from src.api.endpoints.annotate.all.get.models.response import GetNextURLForAllAnnotationResponse from src.api.endpoints.annotate.all.get.queries.core import GetNextURLForAllAnnotationQueryBuilder +from src.api.endpoints.annotate.all.post.models.agency import AnnotationPostAgencyInfo from src.api.endpoints.annotate.all.post.models.name import AnnotationPostNameInfo from src.api.endpoints.annotate.all.post.models.request import AllAnnotationPostInfo from src.core.enums import RecordType @@ -64,7 +65,7 @@ async def test_annotate_all( all_annotations_post_info=AllAnnotationPostInfo( suggested_status=URLType.DATA_SOURCE, record_type=RecordType.ACCIDENT_REPORTS, - agency_ids=[agency_id], + agency_info=AnnotationPostAgencyInfo(agency_ids=[agency_id]), location_ids=[ california.location_id, pennsylvania.location_id, @@ -85,7 +86,7 @@ async def test_annotate_all( all_annotations_post_info=AllAnnotationPostInfo( suggested_status=URLType.NOT_RELEVANT, location_ids=[], - agency_ids=[], + agency_info=AnnotationPostAgencyInfo(agency_ids=[]), name_info=AnnotationPostNameInfo( existing_name_id=setup_info_2.name_suggestion_id ) diff --git a/tests/automated/integration/api/annotate/all/test_new_agency.py b/tests/automated/integration/api/annotate/all/test_new_agency.py new file mode 100644 index 00000000..7a07b3e8 --- /dev/null +++ b/tests/automated/integration/api/annotate/all/test_new_agency.py @@ -0,0 +1,64 @@ +import pytest + +from src.api.endpoints.annotate.all.post.models.agency import AnnotationPostAgencyInfo, \ + AnnotationNewAgencySuggestionInfo +from src.api.endpoints.annotate.all.post.models.name import AnnotationPostNameInfo +from src.api.endpoints.annotate.all.post.models.request import AllAnnotationPostInfo +from src.core.enums import RecordType +from src.db.models.impl.agency.enums import JurisdictionType, AgencyType +from src.db.models.impl.agency.suggestion.sqlalchemy import NewAgencySuggestion +from src.db.models.impl.flag.url_validated.enums import URLType +from tests.helpers.data_creator.models.creation_info.us_state import USStateCreationInfo +from tests.helpers.setup.final_review.core import setup_for_get_next_url_for_final_review +from tests.helpers.setup.final_review.model import FinalReviewSetupInfo + + +@pytest.mark.asyncio +async def test_add_new_agency( + api_test_helper, + pennsylvania: USStateCreationInfo, +): + """ + Test the process for adding a new agency + Confirm a new agency suggestion is successfully added in the database. + """ + ath = api_test_helper + adb_client = ath.adb_client() + + setup_info_1: FinalReviewSetupInfo = await setup_for_get_next_url_for_final_review( + db_data_creator=ath.db_data_creator, + include_user_annotations=True + ) + url_mapping_1 = setup_info_1.url_mapping + + post_response_1 = await ath.request_validator.post_all_annotations_and_get_next( + url_id=url_mapping_1.url_id, + all_annotations_post_info=AllAnnotationPostInfo( + suggested_status=URLType.DATA_SOURCE, + record_type=RecordType.ACCIDENT_REPORTS, + agency_info=AnnotationPostAgencyInfo( + new_agency_suggestion=AnnotationNewAgencySuggestionInfo( + name="New Agency", + location_id=pennsylvania.location_id, + jurisdiction_type=JurisdictionType.STATE, + agency_type=AgencyType.LAW_ENFORCEMENT, + ) + ), + location_ids=[ + pennsylvania.location_id, + ], + name_info=AnnotationPostNameInfo( + new_name="New Name" + ) + ) + ) + + # Check for existence of new agency suggestion + + suggestions: list[NewAgencySuggestion] = await adb_client.get_all(NewAgencySuggestion) + assert len(suggestions) == 1 + suggestion: NewAgencySuggestion = suggestions[0] + assert suggestion.name == "New Agency" + assert suggestion.location_id == pennsylvania.location_id + assert suggestion.jurisdiction_type == JurisdictionType.STATE + assert suggestion.agency_type == AgencyType.LAW_ENFORCEMENT \ No newline at end of file diff --git a/tests/automated/integration/api/annotate/all/test_post_batch_filtering.py b/tests/automated/integration/api/annotate/all/test_post_batch_filtering.py index bfeccc6b..fc34273f 100644 --- a/tests/automated/integration/api/annotate/all/test_post_batch_filtering.py +++ b/tests/automated/integration/api/annotate/all/test_post_batch_filtering.py @@ -1,5 +1,6 @@ import pytest +from src.api.endpoints.annotate.all.post.models.agency import AnnotationPostAgencyInfo from src.api.endpoints.annotate.all.post.models.request import AllAnnotationPostInfo from src.db.models.impl.flag.url_validated.enums import URLType from tests.helpers.setup.final_review.core import setup_for_get_next_url_for_final_review @@ -31,7 +32,7 @@ async def test_annotate_all_post_batch_filtering(api_test_helper): all_annotations_post_info=AllAnnotationPostInfo( suggested_status=URLType.NOT_RELEVANT, location_ids=[], - agency_ids=[] + agency_info=AnnotationPostAgencyInfo(agency_ids=[]) ) ) diff --git a/tests/automated/integration/api/annotate/all/test_validation_error.py b/tests/automated/integration/api/annotate/all/test_validation_error.py index 9c6e244b..d50eca2f 100644 --- a/tests/automated/integration/api/annotate/all/test_validation_error.py +++ b/tests/automated/integration/api/annotate/all/test_validation_error.py @@ -1,5 +1,6 @@ import pytest +from src.api.endpoints.annotate.all.post.models.agency import AnnotationPostAgencyInfo from src.api.endpoints.annotate.all.post.models.request import AllAnnotationPostInfo from src.core.enums import RecordType from src.core.exceptions import FailedValidationException @@ -25,6 +26,6 @@ async def test_annotate_all_validation_error(api_test_helper): suggested_status=URLType.NOT_RELEVANT, record_type=RecordType.ACCIDENT_REPORTS, location_ids=[], - agency_ids=[] + agency_info=AnnotationPostAgencyInfo(agency_ids=[]) ) ) diff --git a/tests/automated/unit/api/test_all_annotation_post_info.py b/tests/automated/unit/api/test_all_annotation_post_info.py index 4649c655..b19eb1b8 100644 --- a/tests/automated/unit/api/test_all_annotation_post_info.py +++ b/tests/automated/unit/api/test_all_annotation_post_info.py @@ -1,6 +1,7 @@ import pytest from pydantic import BaseModel +from src.api.endpoints.annotate.all.post.models.agency import AnnotationPostAgencyInfo from src.api.endpoints.annotate.all.post.models.request import AllAnnotationPostInfo from src.core.enums import RecordType from src.core.exceptions import FailedValidationException @@ -94,13 +95,13 @@ def test_all_annotation_post_info( AllAnnotationPostInfo( suggested_status=params.suggested_status, record_type=params.record_type, - agency_ids=params.agency_ids, + agency_info=AnnotationPostAgencyInfo(agency_ids=params.agency_ids), location_ids=params.location_ids ) else: AllAnnotationPostInfo( suggested_status=params.suggested_status, record_type=params.record_type, - agency_ids=params.agency_ids, + agency_info=AnnotationPostAgencyInfo(agency_ids=params.agency_ids), location_ids=params.location_ids ) \ No newline at end of file