def testSetViolationOnViolationsResponse(self):
"""Tests that the hydrated StateSupervisionViolation is set
on the StateSupervisionViolationResponse."""
supervision_violation_response = \
StateSupervisionViolationResponse.new_with_defaults(
supervision_violation_response_id=123,
response_type=
StateSupervisionViolationResponseType.PERMANENT_DECISION
)
supervision_violation = \
StateSupervisionViolation.new_with_defaults(
supervision_violation_id=999,
supervision_violation_responses=
[supervision_violation_response],
supervision_violation_types=[
StateSupervisionViolationTypeEntry.new_with_defaults(
violation_type=StateSupervisionViolationType.TECHNICAL
)]
)
supervision_violations_and_responses = (
{'violations': [supervision_violation],
'violation_responses': [supervision_violation_response]
}
)
expected_violation_response = \
StateSupervisionViolationResponse.new_with_defaults(
supervision_violation_response_id=123,
response_type=
StateSupervisionViolationResponseType.PERMANENT_DECISION,
supervision_violation=supervision_violation
)
test_pipeline = TestPipeline()
output = (test_pipeline
| beam.Create([(12345,
supervision_violations_and_responses)])
| 'Set Supervision Violation on '
'Supervision Violation Response' >>
beam.ParDo(
entity_hydration_utils.SetViolationOnViolationsResponse())
)
assert_that(output, equal_to([(12345, expected_violation_response)]))
test_pipeline.run()
def testSetViolationOnViolationsResponse(self):
"""Tests that the hydrated StateSupervisionViolation is set
on the StateSupervisionViolationResponse."""
supervision_violation_response = (
StateSupervisionViolationResponse.new_with_defaults(
state_code="US_XX",
supervision_violation_response_id=123,
response_type=StateSupervisionViolationResponseType.
PERMANENT_DECISION,
))
supervision_violation = StateSupervisionViolation.new_with_defaults(
state_code="US_XX",
supervision_violation_id=999,
supervision_violation_responses=[supervision_violation_response],
supervision_violation_types=[
StateSupervisionViolationTypeEntry.new_with_defaults(
state_code="US_XX",
violation_type=StateSupervisionViolationType.TECHNICAL,
)
],
)
supervision_violations_and_responses = {
"violations": [supervision_violation],
"violation_responses": [supervision_violation_response],
}
expected_violation_response = (
StateSupervisionViolationResponse.new_with_defaults(
state_code="US_XX",
supervision_violation_response_id=123,
response_type=StateSupervisionViolationResponseType.
PERMANENT_DECISION,
supervision_violation=supervision_violation,
))
test_pipeline = TestPipeline()
output = test_pipeline | beam.Create(
[(12345, supervision_violations_and_responses)]
) | "Set Supervision Violation on " "Supervision Violation Response" >> beam.ParDo(
entity_hydration_utils.SetViolationOnViolationsResponse())
assert_that(output, equal_to([(12345, expected_violation_response)]))
test_pipeline.run()