1  import pytest
 2  
 3  from mlflow.entities._job_status import JobStatus
 4  from mlflow.protos.jobs_pb2 import JobStatus as ProtoJobStatus
 5  
 6  
 7  @pytest.mark.parametrize(
 8      ("status", "expected_proto"),
 9      [
10          (JobStatus.PENDING, ProtoJobStatus.JOB_STATUS_PENDING),
11          (JobStatus.RUNNING, ProtoJobStatus.JOB_STATUS_IN_PROGRESS),
12          (JobStatus.SUCCEEDED, ProtoJobStatus.JOB_STATUS_COMPLETED),
13          (JobStatus.FAILED, ProtoJobStatus.JOB_STATUS_FAILED),
14          (JobStatus.TIMEOUT, ProtoJobStatus.JOB_STATUS_FAILED),
15          (JobStatus.CANCELED, ProtoJobStatus.JOB_STATUS_CANCELED),
16      ],
17  )
18  def test_job_status_to_proto(status, expected_proto):
19      assert status.to_proto() == expected_proto