wehub-resource-sync
156 lines
4.4 KiB
Python
156 lines
4.4 KiB
Python
from unittest import mock
|
|
|
|
from mlflow.server.jobs.progress import (
|
|
JobTracker,
|
|
NoOpTracker,
|
|
_get_job_tracker,
|
|
_set_job_tracker,
|
|
update_status_details,
|
|
)
|
|
|
|
|
|
def test_job_tracker_writes_to_database():
|
|
job_id = "test-job-123"
|
|
tracker = JobTracker(job_id)
|
|
|
|
with mock.patch("mlflow.server.handlers._get_job_store") as mock_get_store:
|
|
mock_store = mock.Mock()
|
|
mock_get_store.return_value = mock_store
|
|
|
|
tracker.update({"stage": "preprocessing"})
|
|
|
|
mock_get_store.assert_called_once()
|
|
mock_store.update_status_details.assert_called_once_with(job_id, {"stage": "preprocessing"})
|
|
|
|
|
|
def test_job_tracker_writes_status_details_with_metadata():
|
|
job_id = "test-job-456"
|
|
tracker = JobTracker(job_id)
|
|
|
|
with mock.patch("mlflow.server.handlers._get_job_store") as mock_get_store:
|
|
mock_store = mock.Mock()
|
|
mock_get_store.return_value = mock_store
|
|
|
|
tracker.update({"stage": "processing", "progress": "50%", "step": "1"})
|
|
|
|
mock_store.update_status_details.assert_called_once_with(
|
|
job_id, {"stage": "processing", "progress": "50%", "step": "1"}
|
|
)
|
|
|
|
|
|
def test_job_tracker_multiple_updates():
|
|
job_id = "test-job-789"
|
|
tracker = JobTracker(job_id)
|
|
|
|
with mock.patch("mlflow.server.handlers._get_job_store") as mock_get_store:
|
|
mock_store = mock.Mock()
|
|
mock_get_store.return_value = mock_store
|
|
|
|
tracker.update({"stage": "stage1", "key": "value1"})
|
|
tracker.update({"stage": "stage2", "key": "value2"})
|
|
|
|
assert mock_store.update_status_details.call_count == 2
|
|
mock_store.update_status_details.assert_any_call(
|
|
job_id, {"stage": "stage1", "key": "value1"}
|
|
)
|
|
mock_store.update_status_details.assert_any_call(
|
|
job_id, {"stage": "stage2", "key": "value2"}
|
|
)
|
|
|
|
|
|
def test_noop_tracker_does_nothing():
|
|
tracker = NoOpTracker()
|
|
|
|
tracker.update({"stage": "stage1"})
|
|
tracker.update({"stage": "stage2", "key": "value"})
|
|
|
|
|
|
def test_get_job_tracker_returns_noop_by_default():
|
|
_set_job_tracker(None)
|
|
|
|
tracker = _get_job_tracker()
|
|
assert isinstance(tracker, NoOpTracker)
|
|
|
|
|
|
def test_set_and_get_job_tracker():
|
|
job_id = "test-job-abc"
|
|
job_tracker = JobTracker(job_id)
|
|
|
|
_set_job_tracker(job_tracker)
|
|
tracker = _get_job_tracker()
|
|
assert tracker is job_tracker
|
|
|
|
_set_job_tracker(None)
|
|
tracker = _get_job_tracker()
|
|
assert isinstance(tracker, NoOpTracker)
|
|
|
|
|
|
def test_job_tracker_is_global():
|
|
job_id = "test-job-global"
|
|
tracker = JobTracker(job_id)
|
|
|
|
_set_job_tracker(tracker)
|
|
|
|
assert _get_job_tracker() is tracker
|
|
|
|
_set_job_tracker(None)
|
|
|
|
|
|
def test_update_status_details_uses_active_tracker():
|
|
job_id = "test-job-xyz"
|
|
tracker = JobTracker(job_id)
|
|
|
|
with mock.patch("mlflow.server.handlers._get_job_store") as mock_get_store:
|
|
mock_store = mock.Mock()
|
|
mock_get_store.return_value = mock_store
|
|
|
|
_set_job_tracker(tracker)
|
|
update_status_details({"stage": "my_stage", "info": "test"})
|
|
|
|
mock_store.update_status_details.assert_called_once_with(
|
|
job_id, {"stage": "my_stage", "info": "test"}
|
|
)
|
|
|
|
_set_job_tracker(None)
|
|
|
|
|
|
def test_update_status_details_is_noop_without_tracker():
|
|
_set_job_tracker(None)
|
|
|
|
update_status_details({"stage": "stage1"})
|
|
update_status_details({"stage": "stage2", "key": "value"})
|
|
|
|
|
|
def test_update_status_details_with_only_stage():
|
|
job_id = "test-job-stage"
|
|
tracker = JobTracker(job_id)
|
|
|
|
with mock.patch("mlflow.server.handlers._get_job_store") as mock_get_store:
|
|
mock_store = mock.Mock()
|
|
mock_get_store.return_value = mock_store
|
|
|
|
_set_job_tracker(tracker)
|
|
update_status_details({"stage": "my_stage"})
|
|
|
|
mock_store.update_status_details.assert_called_once_with(job_id, {"stage": "my_stage"})
|
|
|
|
_set_job_tracker(None)
|
|
|
|
|
|
def test_update_status_details_with_additional_fields():
|
|
job_id = "test-job-fields"
|
|
tracker = JobTracker(job_id)
|
|
|
|
with mock.patch("mlflow.server.handlers._get_job_store") as mock_get_store:
|
|
mock_store = mock.Mock()
|
|
mock_get_store.return_value = mock_store
|
|
|
|
_set_job_tracker(tracker)
|
|
update_status_details({"stage": "my_stage", "progress": "50%", "details": "processing"})
|
|
|
|
mock_store.update_status_details.assert_called_once_with(
|
|
job_id, {"stage": "my_stage", "progress": "50%", "details": "processing"}
|
|
)
|
|
|
|
_set_job_tracker(None)
|