|
25 | 25 | _is_unassigned_attribute, |
26 | 26 | _calculate_training_progress, |
27 | 27 | _calculate_transition_duration, |
| 28 | + get_mlflow_url, |
28 | 29 | wait |
29 | 30 | ) |
30 | 31 |
|
@@ -489,3 +490,66 @@ def test_wait_metrics_exception_non_jupyter(self, mock_is_jupyter, mock_setup_ml |
489 | 490 |
|
490 | 491 | # Should complete successfully despite metrics exception |
491 | 492 | training_job.refresh.assert_called() |
| 493 | + |
| 494 | + |
| 495 | +class TestGetMlflowUrl: |
| 496 | + """Test cases for get_mlflow_url function.""" |
| 497 | + |
| 498 | + @patch("sagemaker.train.common_utils.mlflow_url_utils.get_presigned_mlflow_experiment_url") |
| 499 | + def test_delegates_to_shared_helper(self, mock_helper): |
| 500 | + """Test that get_mlflow_url extracts config and delegates to shared helper.""" |
| 501 | + mock_helper.return_value = "https://mlflow.example.com/auth?token=abc#/experiments/42" |
| 502 | + |
| 503 | + training_job = MagicMock() |
| 504 | + training_job.mlflow_config.mlflow_resource_arn = "arn:aws:sagemaker:us-west-2:123:mlflow-app/test" |
| 505 | + training_job.mlflow_config.mlflow_experiment_name = "my-experiment" |
| 506 | + |
| 507 | + result = get_mlflow_url(training_job) |
| 508 | + |
| 509 | + mock_helper.assert_called_once_with( |
| 510 | + "arn:aws:sagemaker:us-west-2:123:mlflow-app/test", |
| 511 | + "my-experiment", |
| 512 | + ) |
| 513 | + assert result == "https://mlflow.example.com/auth?token=abc#/experiments/42" |
| 514 | + |
| 515 | + @patch("sagemaker.train.common_utils.trainer_wait.TrainingJob") |
| 516 | + @patch("sagemaker.train.common_utils.mlflow_url_utils.get_presigned_mlflow_experiment_url") |
| 517 | + def test_accepts_job_name_string(self, mock_helper, mock_tj_class): |
| 518 | + """Test that a string job name is resolved via TrainingJob.get().""" |
| 519 | + mock_helper.return_value = "https://mlflow.example.com/auth" |
| 520 | + mock_tj = MagicMock() |
| 521 | + mock_tj.mlflow_config.mlflow_resource_arn = "arn:aws:sagemaker:us-west-2:123:mlflow-app/test" |
| 522 | + mock_tj.mlflow_config.mlflow_experiment_name = None |
| 523 | + mock_tj_class.get.return_value = mock_tj |
| 524 | + |
| 525 | + result = get_mlflow_url("my-training-job") |
| 526 | + |
| 527 | + mock_tj_class.get.assert_called_once_with(training_job_name="my-training-job") |
| 528 | + assert result == "https://mlflow.example.com/auth" |
| 529 | + |
| 530 | + def test_raises_when_no_mlflow_config(self): |
| 531 | + """Test raises ValueError when training job has no mlflow config.""" |
| 532 | + training_job = MagicMock() |
| 533 | + training_job.mlflow_config = MockUnassignedAttribute() |
| 534 | + |
| 535 | + with pytest.raises(ValueError, match="does not have MLflow configured"): |
| 536 | + get_mlflow_url(training_job) |
| 537 | + |
| 538 | + def test_raises_when_mlflow_config_missing(self): |
| 539 | + """Test raises ValueError when training job lacks mlflow_config attribute.""" |
| 540 | + training_job = MagicMock(spec=[]) # no attributes |
| 541 | + |
| 542 | + with pytest.raises(ValueError, match="does not have MLflow configured"): |
| 543 | + get_mlflow_url(training_job) |
| 544 | + |
| 545 | + @patch("sagemaker.train.common_utils.mlflow_url_utils.get_presigned_mlflow_experiment_url") |
| 546 | + def test_raises_when_helper_returns_none(self, mock_helper): |
| 547 | + """Test raises ValueError when presigned URL generation fails.""" |
| 548 | + mock_helper.return_value = None |
| 549 | + |
| 550 | + training_job = MagicMock() |
| 551 | + training_job.mlflow_config.mlflow_resource_arn = "arn:aws:sagemaker:us-west-2:123:mlflow-app/test" |
| 552 | + training_job.mlflow_config.mlflow_experiment_name = "exp" |
| 553 | + |
| 554 | + with pytest.raises(ValueError, match="Failed to generate presigned MLflow URL"): |
| 555 | + get_mlflow_url(training_job) |
0 commit comments