Updating unit tests for internal-external method changes

Author: aviruthen

sagemaker-train/tests/unit/train/aws_batch/test_batch_api_helper.py

@@ -17,10 +17,10 @@
 from unittest.mock import Mock, patch, MagicMock
 
 from sagemaker.train.aws_batch.batch_api_helper import (
-    submit_service_job,
-    describe_service_job,
-    terminate_service_job,
-    list_service_job,
+    _submit_service_job,
+    _describe_service_job,
+    _terminate_service_job,
+    _list_service_job,
 )
 from .conftest import (
     JOB_NAME,
@@ -56,7 +56,7 @@ def test_submit_service_job_basic(self, mock_get_client):
         mock_client.submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
         mock_get_client.return_value = mock_client
 
-        result = submit_service_job(
+        result = _submit_service_job(
             TRAINING_JOB_PAYLOAD,
             JOB_NAME,
             JOB_QUEUE,
@@ -74,7 +74,7 @@ def test_submit_service_job_with_all_params(self, mock_get_client):
         mock_client.submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
         mock_get_client.return_value = mock_client
 
-        result = submit_service_job(
+        result = _submit_service_job(
             TRAINING_JOB_PAYLOAD,
             JOB_NAME,
             JOB_QUEUE,
@@ -103,7 +103,7 @@ def test_submit_service_job_with_tags(self, mock_get_client):
         payload = TRAINING_JOB_PAYLOAD.copy()
         payload["Tags"] = TRAINING_TAGS
 
-        result = submit_service_job(
+        result = _submit_service_job(
             payload,
             JOB_NAME,
             JOB_QUEUE,
@@ -125,7 +125,7 @@ def test_submit_service_job_payload_serialized(self, mock_get_client):
         mock_client.submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
         mock_get_client.return_value = mock_client
 
-        submit_service_job(
+        _submit_service_job(
             TRAINING_JOB_PAYLOAD,
             JOB_NAME,
             JOB_QUEUE,
@@ -148,7 +148,7 @@ def test_describe_service_job(self, mock_get_client):
         mock_client.describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_RUNNING
         mock_get_client.return_value = mock_client
 
-        result = describe_service_job(JOB_ID)
+        result = _describe_service_job(JOB_ID)
 
         assert result["jobId"] == JOB_ID
         assert result["status"] == "RUNNING"
@@ -165,7 +165,7 @@ def test_terminate_service_job(self, mock_get_client):
         mock_client.terminate_service_job.return_value = {}
         mock_get_client.return_value = mock_client
 
-        result = terminate_service_job(JOB_ID, REASON)
+        result = _terminate_service_job(JOB_ID, REASON)
 
         assert result == {}
         mock_client.terminate_service_job.assert_called_once_with(
@@ -179,7 +179,7 @@ def test_terminate_service_job_default_reason(self, mock_get_client):
         mock_client.terminate_service_job.return_value = {}
         mock_get_client.return_value = mock_client
 
-        terminate_service_job(JOB_ID)
+        _terminate_service_job(JOB_ID)
 
         call_kwargs = mock_client.terminate_service_job.call_args[1]
         assert call_kwargs["jobId"] == JOB_ID
@@ -196,7 +196,7 @@ def test_list_service_job_empty(self, mock_get_client):
         mock_client.list_service_jobs.return_value = LIST_SERVICE_JOB_RESP_EMPTY
         mock_get_client.return_value = mock_client
 
-        gen = list_service_job(JOB_QUEUE)
+        gen = _list_service_job(JOB_QUEUE)
         result = next(gen)
 
         assert result["jobSummaryList"] == []
@@ -209,7 +209,7 @@ def test_list_service_job_with_jobs(self, mock_get_client):
         mock_client.list_service_jobs.return_value = LIST_SERVICE_JOB_RESP_WITH_JOBS
         mock_get_client.return_value = mock_client
 
-        gen = list_service_job(JOB_QUEUE)
+        gen = _list_service_job(JOB_QUEUE)
         result = next(gen)
 
         assert len(result["jobSummaryList"]) == 2
@@ -225,7 +225,7 @@ def test_list_service_job_with_pagination(self, mock_get_client):
         ]
         mock_get_client.return_value = mock_client
 
-        gen = list_service_job(JOB_QUEUE)
+        gen = _list_service_job(JOB_QUEUE)
         first_result = next(gen)
         assert first_result["nextToken"] == NEXT_TOKEN
 
@@ -240,7 +240,7 @@ def test_list_service_job_with_filters(self, mock_get_client):
         mock_get_client.return_value = mock_client
 
         filters = [{"name": "JOB_NAME", "values": [JOB_NAME]}]
-        gen = list_service_job(JOB_QUEUE, filters=filters)
+        gen = _list_service_job(JOB_QUEUE, filters=filters)
         result = next(gen)
 
         call_kwargs = mock_client.list_service_jobs.call_args[1]
@@ -253,7 +253,7 @@ def test_list_service_job_with_status(self, mock_get_client):
         mock_client.list_service_jobs.return_value = LIST_SERVICE_JOB_RESP_WITH_JOBS
         mock_get_client.return_value = mock_client
 
-        gen = list_service_job(JOB_QUEUE, job_status=JOB_STATUS_RUNNING)
+        gen = _list_service_job(JOB_QUEUE, job_status=JOB_STATUS_RUNNING)
         result = next(gen)
 
         call_kwargs = mock_client.list_service_jobs.call_args[1]

sagemaker-train/tests/unit/train/aws_batch/test_training_queue.py

@@ -46,7 +46,7 @@ def test_training_queue_init(self):
 class TestTrainingQueueSubmit:
     """Tests for TrainingQueue.submit method"""
 
-    @patch("sagemaker.train.aws_batch.training_queue.submit_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._submit_service_job")
     def test_submit_model_trainer(self, mock_submit_service_job):
         """Test submit with ModelTrainer"""
         mock_submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
@@ -71,7 +71,7 @@ def test_submit_model_trainer(self, mock_submit_service_job):
         assert queued_job.job_arn == JOB_ARN
         mock_submit_service_job.assert_called_once()
 
-    @patch("sagemaker.train.aws_batch.training_queue.submit_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._submit_service_job")
     def test_submit_with_default_timeout(self, mock_submit_service_job):
         """Test submit uses default timeout when not provided"""
         mock_submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
@@ -96,7 +96,7 @@ def test_submit_with_default_timeout(self, mock_submit_service_job):
         # Timeout should be set to default
         assert call_kwargs[5] is not None
 
-    @patch("sagemaker.train.aws_batch.training_queue.submit_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._submit_service_job")
     def test_submit_with_generated_job_name(self, mock_submit_service_job):
         """Test submit generates job name from payload if not provided"""
         mock_submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
@@ -156,7 +156,7 @@ def test_submit_invalid_training_mode(self):
                 BATCH_TAGS,
             )
 
-    @patch("sagemaker.train.aws_batch.training_queue.submit_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._submit_service_job")
     def test_submit_missing_job_arn_in_response(self, mock_submit_service_job):
         """Test submit raises error when jobArn missing from response"""
         mock_submit_service_job.return_value = {"jobName": JOB_NAME}  # Missing jobArn
@@ -183,7 +183,7 @@ def test_submit_missing_job_arn_in_response(self, mock_submit_service_job):
 class TestTrainingQueueMap:
     """Tests for TrainingQueue.map method"""
 
-    @patch("sagemaker.train.aws_batch.training_queue.submit_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._submit_service_job")
     def test_map_multiple_inputs(self, mock_submit_service_job):
         """Test map submits multiple jobs"""
         mock_submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
@@ -208,7 +208,7 @@ def test_map_multiple_inputs(self, mock_submit_service_job):
         assert len(queued_jobs) == 3
         assert mock_submit_service_job.call_count == 3
 
-    @patch("sagemaker.train.aws_batch.training_queue.submit_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._submit_service_job")
     def test_map_with_job_names(self, mock_submit_service_job):
         """Test map with explicit job names"""
         mock_submit_service_job.return_value = SUBMIT_SERVICE_JOB_RESP
@@ -258,7 +258,7 @@ def test_map_mismatched_job_names_length(self):
 class TestTrainingQueueList:
     """Tests for TrainingQueue.list_jobs method"""
 
-    @patch("sagemaker.train.aws_batch.training_queue.list_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
     def test_list_jobs_default(self, mock_list_service_job):
         """Test list_jobs with default parameters"""
         mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_RESP_WITH_JOBS])
@@ -269,7 +269,7 @@ def test_list_jobs_default(self, mock_list_service_job):
         assert len(jobs) == 2
         assert jobs[0].job_name == JOB_NAME
 
-    @patch("sagemaker.train.aws_batch.training_queue.list_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
     def test_list_jobs_with_name_filter(self, mock_list_service_job):
         """Test list_jobs with job name filter"""
         mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_RESP_WITH_JOBS])
@@ -292,7 +292,7 @@ def test_list_jobs_with_name_filter(self, mock_list_service_job):
         assert filters[0]["name"] == "JOB_NAME", "JOB_NAME filter should be present"
         assert filters[0]["values"] == [JOB_NAME], "Filter values should contain the job name"
 
-    @patch("sagemaker.train.aws_batch.training_queue.list_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
     def test_list_jobs_empty(self, mock_list_service_job):
         """Test list_jobs returns empty list"""
         mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_RESP_EMPTY])
@@ -306,7 +306,7 @@ def test_list_jobs_empty(self, mock_list_service_job):
 class TestTrainingQueueGet:
     """Tests for TrainingQueue.get_job method"""
 
-    @patch("sagemaker.train.aws_batch.training_queue.list_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
     def test_get_job_found(self, mock_list_service_job):
         """Test get_job returns job when found"""
         mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_RESP_WITH_JOBS])
@@ -317,7 +317,7 @@ def test_get_job_found(self, mock_list_service_job):
         assert job.job_name == JOB_NAME
         assert job.job_arn == JOB_ARN
 
-    @patch("sagemaker.train.aws_batch.training_queue.list_service_job")
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
     def test_get_job_not_found(self, mock_list_service_job):
         """Test get_job raises error when job not found"""
         mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_RESP_EMPTY])

sagemaker-train/tests/unit/train/aws_batch/test_training_queued_job.py

@@ -49,7 +49,7 @@ def test_training_queued_job_init(self):
 class TestTrainingQueuedJobDescribe:
     """Tests for TrainingQueuedJob.describe method"""
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_describe(self, mock_describe_service_job):
         """Test describe returns job details"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_RUNNING
@@ -64,7 +64,7 @@ def test_describe(self, mock_describe_service_job):
 class TestTrainingQueuedJobTerminate:
     """Tests for TrainingQueuedJob.terminate method"""
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.terminate_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._terminate_service_job")
     def test_terminate(self, mock_terminate_service_job):
         """Test terminate calls terminate API"""
         mock_terminate_service_job.return_value = {}
@@ -74,7 +74,7 @@ def test_terminate(self, mock_terminate_service_job):
 
         mock_terminate_service_job.assert_called_once_with(JOB_ARN, REASON)
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.terminate_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._terminate_service_job")
     def test_terminate_default_reason(self, mock_terminate_service_job):
         """Test terminate with default reason"""
         mock_terminate_service_job.return_value = {}
@@ -89,7 +89,7 @@ def test_terminate_default_reason(self, mock_terminate_service_job):
 class TestTrainingQueuedJobWait:
     """Tests for TrainingQueuedJob.wait method"""
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_wait_immediate_completion(self, mock_describe_service_job):
         """Test wait returns immediately when job is completed"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_SUCCEEDED
@@ -99,7 +99,7 @@ def test_wait_immediate_completion(self, mock_describe_service_job):
 
         assert result["status"] == JOB_STATUS_SUCCEEDED
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_wait_with_polling(self, mock_describe_service_job):
         """Test wait polls until job completes"""
         mock_describe_service_job.side_effect = [
@@ -114,7 +114,7 @@ def test_wait_with_polling(self, mock_describe_service_job):
         assert result["status"] == JOB_STATUS_SUCCEEDED
         assert mock_describe_service_job.call_count == 3
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_wait_with_timeout(self, mock_describe_service_job):
         """Test wait respects timeout"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_RUNNING
@@ -128,7 +128,7 @@ def test_wait_with_timeout(self, mock_describe_service_job):
         assert end_time - start_time >= 2
         assert result["status"] == JOB_STATUS_RUNNING
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_wait_job_failed(self, mock_describe_service_job):
         """Test wait returns failed status"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_FAILED
@@ -144,7 +144,7 @@ class TestTrainingQueuedJobGetModelTrainer:
 
     @patch("sagemaker.train.aws_batch.training_queued_job._remove_system_tags_in_place_in_model_trainer_object")
     @patch("sagemaker.train.aws_batch.training_queued_job._construct_model_trainer_from_training_job_name")
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_get_model_trainer_success(self, mock_describe_service_job, mock_construct_trainer, mock_remove_tags):
         """Test get_model_trainer returns ModelTrainer when training job created"""
         # Return a real dict (not a mock) so nested dict access works
@@ -160,7 +160,7 @@ def test_get_model_trainer_success(self, mock_describe_service_job, mock_constru
         mock_construct_trainer.assert_called_once()
         mock_remove_tags.assert_called_once_with(mock_trainer)
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_get_model_trainer_no_training_job_pending(self, mock_describe_service_job):
         """Test get_model_trainer raises error when job still pending"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_PENDING
@@ -170,7 +170,7 @@ def test_get_model_trainer_no_training_job_pending(self, mock_describe_service_j
         with pytest.raises(NoTrainingJob):
             queued_job.get_model_trainer()
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_get_model_trainer_no_latest_attempt(self, mock_describe_service_job):
         """Test get_model_trainer raises error when latestAttempt missing"""
         resp = DESCRIBE_SERVICE_JOB_RESP_SUCCEEDED.copy()
@@ -186,7 +186,7 @@ def test_get_model_trainer_no_latest_attempt(self, mock_describe_service_job):
 class TestTrainingQueuedJobResult:
     """Tests for TrainingQueuedJob.result method"""
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_result_success(self, mock_describe_service_job):
         """Test result returns job result when completed"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_SUCCEEDED
@@ -196,7 +196,7 @@ def test_result_success(self, mock_describe_service_job):
 
         assert result["status"] == JOB_STATUS_SUCCEEDED
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_result_timeout(self, mock_describe_service_job):
         """Test result raises TimeoutError when timeout exceeded"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_RUNNING
@@ -210,7 +210,7 @@ def test_result_timeout(self, mock_describe_service_job):
 class TestTrainingQueuedJobAsync:
     """Tests for TrainingQueuedJob async methods"""
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_fetch_job_results_success(self, mock_describe_service_job):
         """Test fetch_job_results returns result when job succeeds"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_SUCCEEDED
@@ -220,7 +220,7 @@ def test_fetch_job_results_success(self, mock_describe_service_job):
 
         assert result["status"] == JOB_STATUS_SUCCEEDED
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_fetch_job_results_failed(self, mock_describe_service_job):
         """Test fetch_job_results raises error when job fails"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_FAILED
@@ -230,7 +230,7 @@ def test_fetch_job_results_failed(self, mock_describe_service_job):
         with pytest.raises(RuntimeError):
             asyncio.run(queued_job.fetch_job_results())
 
-    @patch("sagemaker.train.aws_batch.training_queued_job.describe_service_job")
+    @patch("sagemaker.train.aws_batch.training_queued_job._describe_service_job")
     def test_fetch_job_results_timeout(self, mock_describe_service_job):
         """Test fetch_job_results raises TimeoutError when timeout exceeded"""
         mock_describe_service_job.return_value = DESCRIBE_SERVICE_JOB_RESP_RUNNING