def testClassifyProgramAssignments_us_nd(self):
"""Tests the ClassifyProgramAssignments DoFn."""
fake_person_id = 12345
fake_person = entities.StatePerson.new_with_defaults(
state_code="US_ND",
person_id=fake_person_id,
gender=Gender.MALE,
birthdate=date(1970, 1, 1),
residency_status=ResidencyStatus.PERMANENT,
)
program_assignment = entities.StateProgramAssignment.new_with_defaults(
state_code="US_ND",
program_id="PG3",
program_location_id="XYZ",
referral_date=date(2009, 10, 3),
participation_status=StateProgramAssignmentParticipationStatus.IN_PROGRESS,
start_date=date(2009, 10, 19),
)
assessment = entities.StateAssessment.new_with_defaults(
state_code="US_ND",
assessment_type=StateAssessmentType.ORAS,
assessment_score=33,
assessment_date=date(2009, 7, 10),
)
supervision_period = entities.StateSupervisionPeriod.new_with_defaults(
supervision_period_id=111,
status=StateSupervisionPeriodStatus.TERMINATED,
state_code="US_ND",
start_date=date(2008, 3, 5),
termination_date=date(2010, 3, 1),
supervision_type=StateSupervisionType.PAROLE,
supervision_site="10",
)
supervision_period_to_agent_map = {
"agent_id": 1010,
"person_id": fake_person_id,
"agent_external_id": "OFFICER0009",
"supervision_period_id": supervision_period.supervision_period_id,
}
person_periods = {
"person": [fake_person],
"program_assignments": [program_assignment],
"assessments": [assessment],
"supervision_periods": [supervision_period],
"supervision_period_to_agent_association": [
supervision_period_to_agent_map
],
}
program_events = [
ProgramReferralEvent(
state_code=program_assignment.state_code,
program_id=program_assignment.program_id,
event_date=program_assignment.referral_date,
participation_status=program_assignment.participation_status,
assessment_score=33,
assessment_type=StateAssessmentType.ORAS,
supervision_type=supervision_period.supervision_type,
supervising_officer_external_id="OFFICER0009",
supervising_district_external_id="10",
level_1_supervision_location_external_id="10",
),
ProgramParticipationEvent(
state_code=program_assignment.state_code,
program_id=program_assignment.program_id,
program_location_id=program_assignment.program_location_id,
event_date=date.today(),
is_first_day_in_program=True,
supervision_type=supervision_period.supervision_type,
),
]
correct_output = [(fake_person.person_id, (fake_person, program_events))]
test_pipeline = TestPipeline()
output = (
test_pipeline
| beam.Create([(fake_person_id, person_periods)])
| "Identify Program Events"
>> beam.ParDo(pipeline.ClassifyProgramAssignments())
)
assert_that(output, equal_to(correct_output))
test_pipeline.run()
def test_map_program_combinations_multiple_supervision_types(self):
person = StatePerson.new_with_defaults(person_id=12345,
birthdate=date(1984, 8, 31),
gender=Gender.FEMALE)
race = StatePersonRace.new_with_defaults(state_code='US_ND',
race=Race.WHITE)
person.races = [race]
ethnicity = StatePersonEthnicity.new_with_defaults(
state_code='US_ND', ethnicity=Ethnicity.NOT_HISPANIC)
person.ethnicities = [ethnicity]
event_date = date(2009, 10, 7)
program_events = [
ProgramReferralEvent(state_code='US_ND',
event_date=event_date,
program_id='XXX',
supervision_type=StateSupervisionType.PAROLE,
assessment_score=22,
assessment_type=StateAssessmentType.LSIR,
supervising_officer_external_id='OFFICERZ',
supervising_district_external_id='135'),
ProgramReferralEvent(
state_code='US_ND',
event_date=date(2009, 10, 7),
program_id='XXX',
supervision_type=StateSupervisionType.PROBATION,
assessment_score=22,
assessment_type=StateAssessmentType.LSIR,
supervising_officer_external_id='OFFICERZ',
supervising_district_external_id='135'),
ProgramParticipationEvent(
state_code='US_ND',
event_date=event_date,
program_id='XXX',
program_location_id='YYY',
supervision_type=StateSupervisionType.PAROLE,
),
ProgramParticipationEvent(
state_code='US_ND',
event_date=event_date,
program_id='XXX',
program_location_id='YYY',
supervision_type=StateSupervisionType.PROBATION,
)
]
program_combinations = calculator.map_program_combinations(
person,
program_events,
ALL_METRICS_INCLUSIONS_DICT,
calculation_end_month='2009-10',
calculation_month_count=-1)
expected_combinations_count = expected_metric_combos_count(
program_events, duplicated_months_different_supervision_types=True)
self.assertEqual(expected_combinations_count,
len(program_combinations))
assert all(value == 1 for _combination, value in program_combinations)
parole_combos = 0
probation_combos = 0
for combo, _ in program_combinations:
if combo.get('supervision_type') \
and combo_has_enum_value_for_key(
combo, 'methodology', MetricMethodologyType.PERSON):
# Ensure that all person-based metrics are of type parole
if combo.get('supervision_type') == \
StateSupervisionType.PAROLE:
parole_combos += 1
elif combo.get('supervision_type') == \
StateSupervisionType.PROBATION:
probation_combos += 1
# Assert that there are the same number of parole and probation
# person-based combinations
assert parole_combos == probation_combos
def find_program_participation_events(
program_assignment: StateProgramAssignment,
supervision_periods: List[StateSupervisionPeriod],
) -> List[ProgramParticipationEvent]:
"""Finds instances of actively participating in a program. Produces a ProgramParticipationEvent for each day that
the person was actively participating in the program. If the program_assignment has a participation_status of
IN_PROGRESS and has a set start_date, produces a ProgramParticipationEvent for every day between the start_date and
today. If the program_assignment has a participation_status of DISCHARGED, produces a ProgramParticipationEvent for
every day between the start_date and discharge_date, end date exclusive.
Where possible, identifies what types of supervision the person is on on the date of the participation.
If there are multiple overlapping supervision periods, returns one ProgramParticipationEvent for each supervision
period that overlaps.
Returns a list of ProgramReferralEvents.
"""
program_participation_events: List[ProgramParticipationEvent] = []
state_code = program_assignment.state_code
participation_status = program_assignment.participation_status
if participation_status not in (
StateProgramAssignmentParticipationStatus.IN_PROGRESS,
StateProgramAssignmentParticipationStatus.DISCHARGED,
):
return program_participation_events
start_date = program_assignment.start_date
if not start_date:
return program_participation_events
discharge_date = program_assignment.discharge_date
if discharge_date is None:
if (
participation_status
!= StateProgramAssignmentParticipationStatus.IN_PROGRESS
):
logging.warning(
"StateProgramAssignment with a DISCHARGED status but no discharge_date: %s",
program_assignment,
)
return program_participation_events
# This person is actively participating in this program. Set the discharge_date for tomorrow.
discharge_date = date.today() + relativedelta(days=1)
program_id = (
program_assignment.program_id
if program_assignment.program_id
else EXTERNAL_UNKNOWN_VALUE
)
program_location_id = (
program_assignment.program_location_id
if program_assignment.program_location_id
else EXTERNAL_UNKNOWN_VALUE
)
participation_date = start_date
while participation_date < discharge_date:
overlapping_supervision_periods = (
find_supervision_periods_overlapping_with_date(
participation_date, supervision_periods
)
)
is_first_day_in_program = participation_date == start_date
if overlapping_supervision_periods:
for supervision_period in supervision_periods:
program_participation_events.append(
ProgramParticipationEvent(
state_code=state_code,
event_date=participation_date,
is_first_day_in_program=is_first_day_in_program,
program_id=program_id,
program_location_id=program_location_id,
# TODO(#2891): Use supervision_period_supervision_type
supervision_type=supervision_period.supervision_type,
)
)
else:
program_participation_events.append(
ProgramParticipationEvent(
state_code=state_code,
event_date=participation_date,
is_first_day_in_program=is_first_day_in_program,
program_id=program_id,
program_location_id=program_location_id,
)
)
participation_date = participation_date + relativedelta(days=1)
return program_participation_events
def test_find_program_events(self):
program_assignment = StateProgramAssignment.new_with_defaults(
state_code="US_XX",
program_id="PG3",
referral_date=date(2020, 1, 3),
participation_status=StateProgramAssignmentParticipationStatus.
IN_PROGRESS,
program_location_id="LOCATION X",
start_date=date(2020, 1, 1),
)
assessment = StateAssessment.new_with_defaults(
state_code="US_XX",
assessment_type=StateAssessmentType.ORAS,
assessment_score=33,
assessment_date=date(2019, 7, 10),
)
supervision_period = StateSupervisionPeriod.new_with_defaults(
supervision_period_id=999,
status=StateSupervisionPeriodStatus.TERMINATED,
state_code="US_XX",
start_date=date(2019, 3, 5),
termination_date=date(2020, 10, 1),
supervision_type=StateSupervisionType.PAROLE,
supervision_site="OFFICE_1",
)
program_assignments = [program_assignment]
assessments = [assessment]
supervision_periods = [supervision_period]
program_events = identifier.find_program_events(
program_assignments,
assessments,
supervision_periods,
DEFAULT_SUPERVISION_PERIOD_AGENT_ASSOCIATION_LIST,
)
expected_events = [
ProgramReferralEvent(
state_code=program_assignment.state_code,
event_date=program_assignment.referral_date,
program_id=program_assignment.program_id,
supervision_type=supervision_period.supervision_type,
participation_status=program_assignment.participation_status,
assessment_score=assessment.assessment_score,
assessment_type=assessment.assessment_type,
supervising_officer_external_id="XXX",
supervising_district_external_id="OFFICE_1",
level_1_supervision_location_external_id="OFFICE_1",
level_2_supervision_location_external_id=None,
),
ProgramParticipationEvent(
state_code=program_assignment.state_code,
event_date=program_assignment.start_date,
is_first_day_in_program=True,
program_id=program_assignment.program_id,
program_location_id=program_assignment.program_location_id,
supervision_type=supervision_period.supervision_type,
),
ProgramParticipationEvent(
state_code=program_assignment.state_code,
event_date=program_assignment.start_date +
relativedelta(days=1),
is_first_day_in_program=False,
program_id=program_assignment.program_id,
program_location_id=program_assignment.program_location_id,
supervision_type=supervision_period.supervision_type,
),
]
self.assertListEqual(program_events, expected_events)
def test_map_program_combinations_multiple_supervision_types(self):
person = StatePerson.new_with_defaults(
state_code="US_ND",
person_id=12345,
birthdate=date(1984, 8, 31),
gender=Gender.FEMALE,
)
race = StatePersonRace.new_with_defaults(state_code="US_ND",
race=Race.WHITE)
person.races = [race]
ethnicity = StatePersonEthnicity.new_with_defaults(
state_code="US_ND", ethnicity=Ethnicity.NOT_HISPANIC)
person.ethnicities = [ethnicity]
event_date = date(2009, 10, 7)
program_events = [
ProgramReferralEvent(
state_code="US_ND",
event_date=event_date,
program_id="XXX",
supervision_type=StateSupervisionType.PAROLE,
assessment_score=22,
assessment_type=StateAssessmentType.LSIR,
supervising_officer_external_id="OFFICERZ",
supervising_district_external_id="135",
),
ProgramReferralEvent(
state_code="US_ND",
event_date=date(2009, 10, 7),
program_id="XXX",
supervision_type=StateSupervisionType.PROBATION,
assessment_score=22,
assessment_type=StateAssessmentType.LSIR,
supervising_officer_external_id="OFFICERZ",
supervising_district_external_id="135",
),
ProgramParticipationEvent(
state_code="US_ND",
event_date=event_date,
program_id="XXX",
program_location_id="YYY",
supervision_type=StateSupervisionType.PAROLE,
),
ProgramParticipationEvent(
state_code="US_ND",
event_date=event_date,
program_id="XXX",
program_location_id="YYY",
supervision_type=StateSupervisionType.PROBATION,
),
]
program_combinations = calculator.map_program_combinations(
person,
program_events,
ALL_METRICS_INCLUSIONS_DICT,
calculation_end_month="2009-10",
calculation_month_count=-1,
person_metadata=_DEFAULT_PERSON_METADATA,
)
expected_combinations_count = expected_metric_combos_count(
program_events)
self.assertEqual(expected_combinations_count,
len(program_combinations))
assert all(value == 1 for _combination, value in program_combinations)
def testClassifyProgramAssignments(self):
"""Tests the ClassifyProgramAssignments DoFn."""
fake_person_id = 12345
fake_person = entities.StatePerson.new_with_defaults(
person_id=fake_person_id,
gender=Gender.MALE,
birthdate=date(1970, 1, 1),
residency_status=ResidencyStatus.PERMANENT)
program_assignment = entities.StateProgramAssignment.new_with_defaults(
state_code='US_CA',
program_id='PG3',
program_location_id='XYZ',
referral_date=date(2009, 10, 3),
participation_status=StateProgramAssignmentParticipationStatus.
IN_PROGRESS,
start_date=date(2009, 10, 19))
assessment = entities.StateAssessment.new_with_defaults(
state_code='US_CA',
assessment_type=StateAssessmentType.ORAS,
assessment_score=33,
assessment_date=date(2009, 7, 10))
supervision_period = \
entities.StateSupervisionPeriod.new_with_defaults(
supervision_period_id=111,
status=StateSupervisionPeriodStatus.TERMINATED,
state_code='UT',
start_date=date(2008, 3, 5),
supervision_type=StateSupervisionType.PAROLE
)
person_periods = {
'person': [fake_person],
'program_assignments': [program_assignment],
'assessments': [assessment],
'supervision_periods': [supervision_period]
}
program_events = [
ProgramReferralEvent(
state_code=program_assignment.state_code,
program_id=program_assignment.program_id,
event_date=program_assignment.referral_date,
participation_status=program_assignment.participation_status,
assessment_score=33,
assessment_type=StateAssessmentType.ORAS,
supervision_type=supervision_period.supervision_type,
supervising_officer_external_id='OFFICER0009',
supervising_district_external_id='10'),
ProgramParticipationEvent(
state_code=program_assignment.state_code,
program_id=program_assignment.program_id,
program_location_id=program_assignment.program_location_id,
event_date=date.today(),
supervision_type=supervision_period.supervision_type)
]
correct_output = [(fake_person, program_events)]
test_pipeline = TestPipeline()
supervision_period_to_agent_map = {
'agent_id': 1010,
'agent_external_id': 'OFFICER0009',
'district_external_id': '10',
'supervision_period_id': supervision_period.supervision_period_id
}
supervision_period_to_agent_associations = (
test_pipeline
| 'Create SupervisionPeriod to Agent table' >> beam.Create(
[supervision_period_to_agent_map]))
supervision_periods_to_agent_associations_as_kv = (
supervision_period_to_agent_associations
| 'Convert SupervisionPeriod to Agent table to KV tuples' >>
beam.ParDo(pipeline.ConvertDictToKVTuple(),
'supervision_period_id'))
output = (test_pipeline
| beam.Create([(fake_person_id, person_periods)])
| 'Identify Program Events' >> beam.ParDo(
pipeline.ClassifyProgramAssignments(),
AsDict(supervision_periods_to_agent_associations_as_kv)))
assert_that(output, equal_to(correct_output))
test_pipeline.run()