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()