def generate_supervision_period(person,
**kwargs) -> schema.StateSupervisionPeriod:
args = {
"state_code": _STATE_CODE,
"status": StateSupervisionPeriodStatus.PRESENT_WITHOUT_INFO.value,
}
args.update(kwargs)
return schema.StateSupervisionPeriod(person=person, **args)
def test_removeSeosFromViolationIds_unexpectedFormat(self):
with pytest.raises(ValueError):
sv = schema.StateSupervisionViolation(external_id='bad_id')
sp = schema.StateSupervisionPeriod(
supervision_violation_entries=[sv])
ss = schema.StateSupervisionSentence(supervision_periods=[sp])
sg = schema.StateSentenceGroup(supervision_sentences=[ss])
p = schema.StatePerson(sentence_groups=[sg])
remove_suffix_from_violation_ids([p])
def _move_violations_onto_supervision_periods_for_sentence(
sentence: Union[StateSupervisionSentence, StateIncarcerationSentence]):
"""Looks at all SupervisionViolations in the provided |sentence|, and attempts to match them to the corresponding
SupervisionPeriod, based on date.
"""
supervision_periods = sentence.supervision_periods
# Get all supervision violations from supervision periods
supervision_violations = get_all_entities_of_cls(supervision_periods, StateSupervisionViolation)
# Clear the links from supervision period to supervision violations. We will
# re-add/update these relationships below.
for supervision_period in supervision_periods:
supervision_period.supervision_violation_entries = []
unmatched_svs = []
non_placeholder_periods = [sp for sp in supervision_periods if not is_placeholder(sp)]
# Match SVs to non_placeholder_periods by date.
for sv in supervision_violations:
matched = False
violation_date = _get_approximate_violation_date(sv)
if violation_date:
for sp in non_placeholder_periods:
sp_end_date = sp.termination_date if sp.termination_date else datetime.date.max
sp_start_date = sp.start_date
if sp_start_date <= violation_date <= sp_end_date:
matched = True
sp.supervision_violation_entries.append(sv)
# Unmatched SVs will be re-added to a placeholder period at the end.
if not matched:
unmatched_svs.append(sv)
# Add unmatched supervision violations to a placeholder period
if unmatched_svs:
placeholder_periods = [sp for sp in supervision_periods if is_placeholder(sp)]
if not placeholder_periods:
# We may hit this case if an entity that has already been committed to the DB has a date updated in a later
# run such that the dates of the existing supervision periods no longer line up with one of the existing
# supervision violations.
logging.info(
'No placeholder supervision periods exist on sentence [%s]([%s]), creating a new placeholder '
'supervision period.',
sentence.external_id, sentence.get_id())
new_placeholder_supervision_period = schema.StateSupervisionPeriod(
state_code=sentence.state_code,
status=StateSupervisionPeriodStatus.PRESENT_WITHOUT_INFO.value,
person=sentence.person
)
placeholder_periods.append(new_placeholder_supervision_period)
sentence.supervision_periods.append(new_placeholder_supervision_period)
placeholder_periods[0].supervision_violation_entries = unmatched_svs
def generate_test_supervision_period(person_id, supervision_violations) -> \
state_schema.StateSupervisionPeriod:
instance = state_schema.StateSupervisionPeriod(
supervision_period_id=4444,
status=StateSupervisionPeriodStatus.EXTERNAL_UNKNOWN.value,
state_code='us_ca',
person_id=person_id,
supervision_violations=supervision_violations,
)
return instance
def generate_test_supervision_period(
person_id, supervision_violations, case_types,
supervision_contacts) -> state_schema.StateSupervisionPeriod:
instance = state_schema.StateSupervisionPeriod(
supervision_period_id=4444,
status=StateSupervisionPeriodStatus.EXTERNAL_UNKNOWN.value,
state_code="us_ca",
person_id=person_id,
case_type_entries=case_types,
supervision_violation_entries=supervision_violations,
supervision_contacts=supervision_contacts,
)
return instance
def test_removeSeosFromViolationIds_unexpectedFormat(self) -> None:
with pytest.raises(ValueError) as e:
sv = schema.StateSupervisionViolation(external_id="bad_id")
sp = schema.StateSupervisionPeriod(
supervision_violation_entries=[sv],
status=StateSupervisionPeriodStatus.PRESENT_WITHOUT_INFO,
)
ss = schema.StateSupervisionSentence(supervision_periods=[sp])
sg = schema.StateSentenceGroup(supervision_sentences=[ss])
p = schema.StatePerson(sentence_groups=[sg])
remove_suffix_from_violation_ids([p])
self.assertEqual(
str(e.value),
"Unexpected id format [bad_id] for [StateSupervisionViolation(external_id=bad_id)]",
)
def test_removeSeosFromViolationIds(self):
svr = schema.StateSupervisionViolationResponse(
external_id='DOC-CYC-VSN1-SEO-FSO')
sv = schema.StateSupervisionViolation(
external_id='DOC-CYC-VSN1-SEO-FSO',
supervision_violation_responses=[svr])
svr_2 = schema.StateSupervisionViolationResponse(
external_id='DOC-CYC-VSN1-SEO-FSO')
sv_2 = schema.StateSupervisionViolation(
external_id='DOC-CYC-VSN1-SEO-FSO',
supervision_violation_responses=[svr_2])
sp = schema.StateSupervisionPeriod(
supervision_violation_entries=[sv, sv_2])
ss = schema.StateSupervisionSentence(supervision_periods=[sp])
sg = schema.StateSentenceGroup(supervision_sentences=[ss])
p = schema.StatePerson(sentence_groups=[sg])
expected_svr = StateSupervisionViolationResponse.new_with_defaults(
external_id='DOC-CYC-VSN1')
expected_sv = StateSupervisionViolation.new_with_defaults(
external_id='DOC-CYC-VSN1',
supervision_violation_responses=[expected_svr])
expected_svr_2 = attr.evolve(expected_svr)
expected_sv_2 = attr.evolve(
expected_sv,
supervision_violation_responses=[expected_svr_2])
expected_sp = StateSupervisionPeriod.new_with_defaults(
supervision_violation_entries=[expected_sv, expected_sv_2])
expected_ss = StateSupervisionSentence.new_with_defaults(
supervision_periods=[expected_sp])
expected_sg = StateSentenceGroup.new_with_defaults(
supervision_sentences=[expected_ss])
expected_p = StatePerson.new_with_defaults(
sentence_groups=[expected_sg])
remove_suffix_from_violation_ids([p])
self.assertEqual(expected_p, self.to_entity(p))
def build_data_dict(fake_person_id: int, fake_supervision_period_id: int):
"""Builds a data_dict for a basic run of the pipeline."""
fake_person = schema.StatePerson(
person_id=fake_person_id,
gender=Gender.MALE,
birthdate=date(1970, 1, 1),
residency_status=ResidencyStatus.PERMANENT)
persons_data = [normalized_database_base_dict(fake_person)]
race_1 = schema.StatePersonRace(person_race_id=111,
state_code='CA',
race=Race.BLACK,
person_id=fake_person_id)
race_2 = schema.StatePersonRace(person_race_id=111,
state_code='ND',
race=Race.WHITE,
person_id=fake_person_id)
races_data = normalized_database_base_dict_list([race_1, race_2])
ethnicity = schema.StatePersonEthnicity(person_ethnicity_id=111,
state_code='CA',
ethnicity=Ethnicity.HISPANIC,
person_id=fake_person_id)
ethnicity_data = normalized_database_base_dict_list([ethnicity])
program_assignment = schema.StateProgramAssignment(
state_code='CA',
program_assignment_id=123,
referral_date=date(2015, 5, 10),
person_id=fake_person_id)
assessment = schema.StateAssessment(assessment_id=298374,
assessment_date=date(2015, 3, 19),
assessment_type='LSIR',
person_id=fake_person_id)
supervision_period = schema.StateSupervisionPeriod(
supervision_period_id=fake_supervision_period_id,
state_code='CA',
county_code='124',
start_date=date(2015, 3, 14),
termination_date=date(2016, 12, 29),
supervision_type=StateSupervisionType.PROBATION,
person_id=fake_person_id)
program_assignment_data = [
normalized_database_base_dict(program_assignment)
]
assessment_data = [normalized_database_base_dict(assessment)]
supervision_periods_data = [
normalized_database_base_dict(supervision_period)
]
supervision_violation_response = \
database_test_utils.generate_test_supervision_violation_response(
fake_person_id)
supervision_violation_response_data = [
normalized_database_base_dict(supervision_violation_response)
]
data_dict = {
schema.StatePerson.__tablename__: persons_data,
schema.StatePersonRace.__tablename__: races_data,
schema.StatePersonEthnicity.__tablename__: ethnicity_data,
schema.StateSupervisionViolationResponse.__tablename__:
supervision_violation_response_data,
schema.StateSupervisionPeriod.__tablename__:
supervision_periods_data,
schema.StateProgramAssignment.__tablename__:
program_assignment_data,
schema.StateAssessment.__tablename__: assessment_data,
schema.StatePersonExternalId.__tablename__: [],
schema.StatePersonAlias.__tablename__: [],
schema.StateSentenceGroup.__tablename__: [],
}
return data_dict
def testProgramPipelineNoReferrals(self):
"""Tests the program pipeline where one person does not have any
program assignment entities."""
fake_person_id = 12345
fake_person_id_2 = 9876
fake_person = schema.StatePerson(
person_id=fake_person_id,
gender=Gender.MALE,
birthdate=date(1970, 1, 1),
residency_status=ResidencyStatus.PERMANENT)
fake_person_2 = schema.StatePerson(
person_id=fake_person_id_2,
gender=Gender.MALE,
birthdate=date(1974, 3, 12),
residency_status=ResidencyStatus.PERMANENT)
persons_data = normalized_database_base_dict_list(
[fake_person, fake_person_2])
race_1 = schema.StatePersonRace(person_race_id=111,
state_code='CA',
race=Race.BLACK,
person_id=fake_person_id)
race_2 = schema.StatePersonRace(person_race_id=111,
state_code='ND',
race=Race.WHITE,
person_id=fake_person_id)
races_data = normalized_database_base_dict_list([race_1, race_2])
ethnicity = schema.StatePersonEthnicity(person_ethnicity_id=111,
state_code='CA',
ethnicity=Ethnicity.HISPANIC,
person_id=fake_person_id)
ethnicity_data = normalized_database_base_dict_list([ethnicity])
# Program assignment for a different person
program_assignment = schema.StateProgramAssignment(
state_code='CA',
program_assignment_id=123,
referral_date=date(2015, 5, 10),
person_id=fake_person_id_2)
assessment = schema.StateAssessment(assessment_id=298374,
assessment_date=date(2015, 3, 19),
assessment_type='LSIR',
person_id=fake_person_id)
supervision_period = schema.StateSupervisionPeriod(
supervision_period_id=1111,
state_code='CA',
county_code='124',
start_date=date(2015, 3, 14),
termination_date=date(2016, 12, 29),
supervision_type=StateSupervisionType.PROBATION,
person_id=fake_person_id)
program_assignment_data = [
normalized_database_base_dict(program_assignment)
]
assessment_data = [normalized_database_base_dict(assessment)]
supervision_periods_data = [
normalized_database_base_dict(supervision_period)
]
supervision_violation_response = \
database_test_utils.generate_test_supervision_violation_response(
fake_person_id)
supervision_violation_response_data = [
normalized_database_base_dict(supervision_violation_response)
]
data_dict = {
schema.StatePerson.__tablename__: persons_data,
schema.StatePersonRace.__tablename__: races_data,
schema.StatePersonEthnicity.__tablename__: ethnicity_data,
schema.StateSupervisionViolationResponse.__tablename__:
supervision_violation_response_data,
schema.StateSupervisionPeriod.__tablename__:
supervision_periods_data,
schema.StateProgramAssignment.__tablename__:
program_assignment_data,
schema.StateAssessment.__tablename__: assessment_data,
schema.StatePersonExternalId.__tablename__: [],
schema.StatePersonAlias.__tablename__: [],
schema.StateSentenceGroup.__tablename__: [],
}
dataset = 'recidiviz-123.state'
with patch(
'recidiviz.calculator.pipeline.utils.extractor_utils.ReadFromBigQuery',
self.fake_bq_source_factory.create_fake_bq_source_constructor(
dataset, data_dict)):
self.run_test_pipeline(dataset,
supervision_period.supervision_period_id)
def testProgramPipelineNoReferrals(self):
"""Tests the program pipeline where one person does not have any
program assignment entities."""
fake_person_id = 12345
fake_person_id_2 = 9876
fake_person = schema.StatePerson(
state_code="US_XX",
person_id=fake_person_id,
gender=Gender.MALE,
birthdate=date(1970, 1, 1),
residency_status=ResidencyStatus.PERMANENT,
)
fake_person_2 = schema.StatePerson(
state_code="US_XX",
person_id=fake_person_id_2,
gender=Gender.MALE,
birthdate=date(1974, 3, 12),
residency_status=ResidencyStatus.PERMANENT,
)
persons_data = normalized_database_base_dict_list([fake_person, fake_person_2])
race_1 = schema.StatePersonRace(
person_race_id=111,
state_code="US_XX",
race=Race.BLACK,
person_id=fake_person_id,
)
race_2 = schema.StatePersonRace(
person_race_id=111,
state_code="US_XX",
race=Race.WHITE,
person_id=fake_person_id,
)
races_data = normalized_database_base_dict_list([race_1, race_2])
ethnicity = schema.StatePersonEthnicity(
person_ethnicity_id=111,
state_code="US_XX",
ethnicity=Ethnicity.HISPANIC,
person_id=fake_person_id,
)
ethnicity_data = normalized_database_base_dict_list([ethnicity])
# Program assignment for a different person
program_assignment = schema.StateProgramAssignment(
state_code="US_XX",
program_assignment_id=123,
referral_date=date(2015, 5, 10),
person_id=fake_person_id_2,
participation_status=StateProgramAssignmentParticipationStatus.DENIED,
)
assessment = schema.StateAssessment(
assessment_id=298374,
state_code="US_XX",
assessment_date=date(2015, 3, 19),
assessment_type="LSIR",
person_id=fake_person_id,
)
supervision_period = schema.StateSupervisionPeriod(
supervision_period_id=1111,
state_code="US_XX",
county_code="124",
start_date=date(2015, 3, 14),
termination_date=date(2016, 12, 29),
supervision_type=StateSupervisionType.PROBATION,
person_id=fake_person_id,
status=StateSupervisionPeriodStatus.PRESENT_WITHOUT_INFO,
)
program_assignment_data = [normalized_database_base_dict(program_assignment)]
assessment_data = [normalized_database_base_dict(assessment)]
supervision_periods_data = [normalized_database_base_dict(supervision_period)]
supervision_violation_response = (
database_test_utils.generate_test_supervision_violation_response(
fake_person_id
)
)
supervision_violation_response_data = [
normalized_database_base_dict(supervision_violation_response)
]
supervision_period_to_agent_data = [
{
"agent_id": 1010,
"person_id": fake_person_id,
"state_code": "US_XX",
"agent_external_id": "OFFICER0009",
"supervision_period_id": supervision_period.supervision_period_id,
}
]
state_race_ethnicity_population_count_data = [
{
"state_code": "US_XX",
"race_or_ethnicity": "BLACK",
"population_count": 1,
"representation_priority": 1,
}
]
data_dict = {
schema.StatePerson.__tablename__: persons_data,
schema.StatePersonRace.__tablename__: races_data,
schema.StatePersonEthnicity.__tablename__: ethnicity_data,
schema.StateSupervisionViolationResponse.__tablename__: supervision_violation_response_data,
schema.StateSupervisionPeriod.__tablename__: supervision_periods_data,
schema.StateProgramAssignment.__tablename__: program_assignment_data,
schema.StateAssessment.__tablename__: assessment_data,
schema.StatePersonExternalId.__tablename__: [],
schema.StatePersonAlias.__tablename__: [],
schema.StateSentenceGroup.__tablename__: [],
"supervision_period_to_agent_association": supervision_period_to_agent_data,
"state_race_ethnicity_population_counts": state_race_ethnicity_population_count_data,
}
dataset = "recidiviz-123.state"
self.run_test_pipeline(dataset, data_dict)
def build_data_dict(fake_person_id: int, fake_supervision_period_id: int):
"""Builds a data_dict for a basic run of the pipeline."""
fake_person = schema.StatePerson(
state_code="US_XX",
person_id=fake_person_id,
gender=Gender.MALE,
birthdate=date(1970, 1, 1),
residency_status=ResidencyStatus.PERMANENT,
)
persons_data = [normalized_database_base_dict(fake_person)]
race_1 = schema.StatePersonRace(
person_race_id=111,
state_code="US_XX",
race=Race.BLACK,
person_id=fake_person_id,
)
race_2 = schema.StatePersonRace(
person_race_id=111,
state_code="US_XX",
race=Race.WHITE,
person_id=fake_person_id,
)
races_data = normalized_database_base_dict_list([race_1, race_2])
ethnicity = schema.StatePersonEthnicity(
person_ethnicity_id=111,
state_code="US_XX",
ethnicity=Ethnicity.HISPANIC,
person_id=fake_person_id,
)
ethnicity_data = normalized_database_base_dict_list([ethnicity])
program_assignment = schema.StateProgramAssignment(
state_code="US_XX",
program_assignment_id=123,
referral_date=date(2015, 5, 10),
person_id=fake_person_id,
participation_status=StateProgramAssignmentParticipationStatus.IN_PROGRESS,
)
assessment = schema.StateAssessment(
assessment_id=298374,
state_code="US_XX",
assessment_date=date(2015, 3, 19),
assessment_type="LSIR",
person_id=fake_person_id,
)
supervision_period = schema.StateSupervisionPeriod(
supervision_period_id=fake_supervision_period_id,
state_code="US_XX",
county_code="124",
start_date=date(2015, 3, 14),
termination_date=date(2016, 12, 29),
supervision_type=StateSupervisionType.PROBATION,
person_id=fake_person_id,
status=StateSupervisionPeriodStatus.PRESENT_WITHOUT_INFO,
)
program_assignment_data = [normalized_database_base_dict(program_assignment)]
assessment_data = [normalized_database_base_dict(assessment)]
supervision_periods_data = [normalized_database_base_dict(supervision_period)]
supervision_violation_response = (
database_test_utils.generate_test_supervision_violation_response(
fake_person_id
)
)
supervision_violation_response_data = [
normalized_database_base_dict(supervision_violation_response)
]
supervision_period_to_agent_data = [
{
"agent_id": 1010,
"person_id": fake_person_id,
"state_code": "US_XX",
"agent_external_id": "OFFICER0009",
"supervision_period_id": fake_supervision_period_id,
}
]
state_race_ethnicity_population_count_data = [
{
"state_code": "US_XX",
"race_or_ethnicity": "BLACK",
"population_count": 1,
"representation_priority": 1,
}
]
data_dict = {
schema.StatePerson.__tablename__: persons_data,
schema.StatePersonRace.__tablename__: races_data,
schema.StatePersonEthnicity.__tablename__: ethnicity_data,
schema.StateSupervisionViolationResponse.__tablename__: supervision_violation_response_data,
schema.StateSupervisionPeriod.__tablename__: supervision_periods_data,
schema.StateProgramAssignment.__tablename__: program_assignment_data,
schema.StateAssessment.__tablename__: assessment_data,
schema.StatePersonExternalId.__tablename__: [],
schema.StatePersonAlias.__tablename__: [],
schema.StateSentenceGroup.__tablename__: [],
"supervision_period_to_agent_association": supervision_period_to_agent_data,
"state_race_ethnicity_population_counts": state_race_ethnicity_population_count_data,
}
return data_dict
def test_removeSeosFromViolationIds(self) -> None:
svr = schema.StateSupervisionViolationResponse(
state_code=_STATE_CODE, external_id="DOC-CYC-VSN1-SEO-FSO")
sv = schema.StateSupervisionViolation(
state_code=_STATE_CODE,
external_id="DOC-CYC-VSN1-SEO-FSO",
supervision_violation_responses=[svr],
)
svr_2 = schema.StateSupervisionViolationResponse(
state_code=_STATE_CODE, external_id="DOC-CYC-VSN1-SEO-FSO")
sv_2 = schema.StateSupervisionViolation(
state_code=_STATE_CODE,
external_id="DOC-CYC-VSN1-SEO-FSO",
supervision_violation_responses=[svr_2],
)
sp = schema.StateSupervisionPeriod(
state_code=_STATE_CODE,
supervision_violation_entries=[sv, sv_2],
status=StateSupervisionPeriodStatus.PRESENT_WITHOUT_INFO,
)
ss = schema.StateSupervisionSentence(
state_code=_STATE_CODE,
supervision_periods=[sp],
status=StateSentenceStatus.PRESENT_WITHOUT_INFO,
)
sg = schema.StateSentenceGroup(
state_code=_STATE_CODE,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO,
supervision_sentences=[ss],
)
p = schema.StatePerson(state_code=_STATE_CODE, sentence_groups=[sg])
expected_svr = StateSupervisionViolationResponse.new_with_defaults(
state_code=_STATE_CODE, external_id="DOC-CYC-VSN1")
expected_sv = StateSupervisionViolation.new_with_defaults(
state_code=_STATE_CODE,
external_id="DOC-CYC-VSN1",
supervision_violation_responses=[expected_svr],
)
expected_svr_2 = attr.evolve(expected_svr)
expected_sv_2 = attr.evolve(
expected_sv, supervision_violation_responses=[expected_svr_2])
expected_sp = StateSupervisionPeriod.new_with_defaults(
state_code=_STATE_CODE,
supervision_violation_entries=[expected_sv, expected_sv_2],
status=StateSupervisionPeriodStatus.PRESENT_WITHOUT_INFO,
)
expected_ss = StateSupervisionSentence.new_with_defaults(
state_code=_STATE_CODE,
supervision_periods=[expected_sp],
status=StateSentenceStatus.PRESENT_WITHOUT_INFO,
)
expected_sg = StateSentenceGroup.new_with_defaults(
state_code=_STATE_CODE,
status=StateSentenceStatus.PRESENT_WITHOUT_INFO,
supervision_sentences=[expected_ss],
)
expected_p = StatePerson.new_with_defaults(
state_code=_STATE_CODE, sentence_groups=[expected_sg])
remove_suffix_from_violation_ids([p])
self.assertEqual(expected_p, self.to_entity(p))
def testProgramPipeline(self):
"""Tests the program pipeline."""
fake_person_id = 12345
fake_person = schema.StatePerson(
person_id=fake_person_id,
gender=Gender.MALE,
birthdate=date(1970, 1, 1),
residency_status=ResidencyStatus.PERMANENT)
persons_data = [normalized_database_base_dict(fake_person)]
race_1 = schema.StatePersonRace(person_race_id=111,
state_code='CA',
race=Race.BLACK,
person_id=fake_person_id)
race_2 = schema.StatePersonRace(person_race_id=111,
state_code='ND',
race=Race.WHITE,
person_id=fake_person_id)
races_data = normalized_database_base_dict_list([race_1, race_2])
ethnicity = schema.StatePersonEthnicity(person_ethnicity_id=111,
state_code='CA',
ethnicity=Ethnicity.HISPANIC,
person_id=fake_person_id)
ethnicity_data = normalized_database_base_dict_list([ethnicity])
program_assignment = schema.StateProgramAssignment(
program_assignment_id=123,
referral_date=date(2015, 5, 10),
person_id=fake_person_id)
assessment = schema.StateAssessment(assessment_id=298374,
assessment_date=date(2015, 3, 19),
assessment_type='LSIR',
person_id=fake_person_id)
supervision_period = schema.StateSupervisionPeriod(
supervision_period_id=1111,
state_code='CA',
county_code='124',
start_date=date(2015, 3, 14),
termination_date=date(2016, 12, 29),
supervision_type=StateSupervisionType.PROBATION,
person_id=fake_person_id)
program_assignment_data = [
normalized_database_base_dict(program_assignment)
]
assessment_data = [normalized_database_base_dict(assessment)]
supervision_periods_data = [
normalized_database_base_dict(supervision_period)
]
supervision_violation_response = \
database_test_utils.generate_test_supervision_violation_response(
fake_person_id)
supervision_violation_response_data = [
normalized_database_base_dict(supervision_violation_response)
]
data_dict = {
schema.StatePerson.__tablename__: persons_data,
schema.StatePersonRace.__tablename__: races_data,
schema.StatePersonEthnicity.__tablename__: ethnicity_data,
schema.StateSupervisionViolationResponse.__tablename__:
supervision_violation_response_data,
schema.StateSupervisionPeriod.__tablename__:
supervision_periods_data,
schema.StateProgramAssignment.__tablename__:
program_assignment_data,
schema.StateAssessment.__tablename__: assessment_data
}
test_pipeline = TestPipeline()
# Get StatePersons
persons = (test_pipeline
| 'Load Persons' >> extractor_utils.BuildRootEntity(
dataset=None,
data_dict=data_dict,
root_schema_class=schema.StatePerson,
root_entity_class=entities.StatePerson,
unifying_id_field='person_id',
build_related_entities=True))
# Get StateProgramAssignments
program_assignments = (
test_pipeline
| 'Load Program Assignments' >> extractor_utils.BuildRootEntity(
dataset=None,
data_dict=data_dict,
root_schema_class=schema.StateProgramAssignment,
root_entity_class=entities.StateProgramAssignment,
unifying_id_field='person_id',
build_related_entities=True))
# Get StateAssessments
assessments = (test_pipeline
| 'Load Assessments' >> extractor_utils.BuildRootEntity(
dataset=None,
data_dict=data_dict,
root_schema_class=schema.StateAssessment,
root_entity_class=entities.StateAssessment,
unifying_id_field='person_id',
build_related_entities=False))
# Get StateSupervisionPeriods
supervision_periods = (
test_pipeline
| 'Load SupervisionPeriods' >> extractor_utils.BuildRootEntity(
dataset=None,
data_dict=data_dict,
root_schema_class=schema.StateSupervisionPeriod,
root_entity_class=entities.StateSupervisionPeriod,
unifying_id_field='person_id',
build_related_entities=False))
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_period_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'))
# Group each StatePerson with their other entities
persons_entities = ({
'person': persons,
'program_assignments': program_assignments,
'assessments': assessments,
'supervision_periods': supervision_periods
}
|
'Group StatePerson to StateProgramAssignments and'
>> beam.CoGroupByKey())
# Identify ProgramEvents from the StatePerson's
# StateProgramAssignments
person_program_events = (
persons_entities
| beam.ParDo(
pipeline.ClassifyProgramAssignments(),
AsDict(supervision_period_to_agent_associations_as_kv)))
# Get pipeline job details for accessing job_id
all_pipeline_options = PipelineOptions().get_all_options()
# Add timestamp for local jobs
job_timestamp = datetime.datetime.now().strftime(
'%Y-%m-%d_%H_%M_%S.%f')
all_pipeline_options['job_timestamp'] = job_timestamp
# Get program metrics
program_metrics = (person_program_events
|
'Get Program Metrics' >> pipeline.GetProgramMetrics(
pipeline_options=all_pipeline_options,
inclusions=ALL_INCLUSIONS_DICT,
calculation_month_limit=-1))
assert_that(program_metrics, AssertMatchers.validate_pipeline_test())
test_pipeline.run()
