test(train): patch mocks at the lookup site so unit tests stop calling AWS

[mohamedzeidan2021]

Problem

Five tests in sagemaker-train's unit suite issue real AWS requests — STS GetCallerIdentity, IAM, SageMaker DescribeHubContent, and S3. They pass in CI only because the build role carries credentials and permissions, so the mocks are silently ineffective rather than visibly broken.

Reproduce with deliberately invalid credentials (note AWS_PROFILE must be unset, not empty — jumpstart/constants.py:114 builds a boto3.Session() at import time):

AWS_ACCESS_KEY_ID=AKIAIOSFODNN7INVALID AWS_SECRET_ACCESS_KEY=... AWS_SESSION_TOKEN=invalid \
AWS_EC2_METADATA_DISABLED=true AWS_CONFIG_FILE=/dev/null AWS_SHARED_CREDENTIALS_FILE=/dev/null \
python -m pytest tests/unit/train/test_sft_trainer.py \
  tests/unit/train/evaluate/test_mtrl_evaluator_handshake.py \
  tests/unit/train/test_trainer_recipe_integration.py
→ 5 failed, 97 passed

Root cause — three shapes

1. Patch targeting the defining module instead of the lookup site.

base_trainer.py:39-47 does from ...finetune_utils import get_recipe_s3_uri, get_training_image and :64 does from sagemaker.train.defaults import TrainDefaults, all at module level. Patching those names where they are defined never rebinds the reference base_trainer already holds, so the real functions ran: get_recipe_s3_uri called DescribeHubContent, and TrainDefaults returned a real Session, so download_file issued a live HeadObject.

Only two of the four finetune_utils patches in that block were affected. _render_recipe_placeholders and _get_smtj_override_spec are imported inside _train_serverful_smtj (base_trainer.py:1060-1064), so that import re-executes while the patch is active and finetune_utils is the correct target for them. They are deliberately left unchanged.

2. Missing mock. base_evaluator.py:20-23 imports verify_evaluation_caller_permissions in the same statement as resolve_and_validate_role, but only the latter was ever patched — so _get_aws_execution_context (base_evaluator.py:827) reached STS and IAM.

3. Incomplete mock. test_model_trainer_no_from_recipe_raises already patched TrainDefaults.get_role for exactly this reason, but STS was reached by a second route: model_post_init builds a default OutputDataConfig, which calls session.default_bucket(). The session is stubbed rather than get_output_data_config, so the defaulting logic stays under test. spec=Session is needed because ModelTrainer validates the field as a real Session.

Verification

Measured with invalid credentials and a transport-level recorder wrapping botocore.httpsession.URLLib3Session.send + requests.adapters.HTTPAdapter.send, so calls swallowed by try/except are counted too.

Check	Before	After
The three files	5 failed, 97 passed	102 passed, 0 failed
Full sagemaker-train unit suite	19 failed, 3446 passed	14 failed, 3451 passed
Outbound AWS requests from these files	STS, IAM, SageMaker, S3	0

The remaining 14 failures are pre-existing and unrelated — a local botocore that lacks the SageMaker training-session shapes. They are identical before and after.

Coverage

Net −1 line (27402 → 27401), with no real loss:

• 12 lines no longer covered, all in sagemaker-core/helper/session_helper.py — default_bucket(), generate_default_sagemaker_bucket_name(), botocore_resolver(), sts_regional_endpoint(). These were covered only because a unit test genuinely called STS. All 12 are covered by sagemaker-core's own unit suite (verified).
• 11 lines newly covered in base_trainer, base_evaluator, benchmark_evaluator, custom_scorer_evaluator, model_trainer — production paths that previously never executed because the tests aborted partway through on an AWS error.

Mock-drift risk is covered: get_recipe_s3_uri is asserted directly at tests/unit/train/common_utils/test_finetune_utils.py:1563-1583, and verify_evaluation_caller_permissions has 77 tests in sagemaker-core/tests/unit/helper/test_iam_role_resolver.py.

Scope

Test files only — no production code changes.

Note for reviewers

Escapes here are layered: fixing the get_recipe_s3_uri target made DescribeHubContent disappear and revealed a live HeadObject underneath from the still-real TrainDefaults session. If you spot further unmocked paths in these files, that is the shape to look for. The zero-outbound-request measurement is the check that closed it out.

The same class of bug exists more widely — notably sagemaker-serve, where 65 unit tests fail under invalid credentials (23 of them because constructing ModelBuilder(model=Mock()) calls STS via __post_init__). That is intentionally out of scope here and worth separate PRs.

[github-actions]

🤖 AI code review

This is a test-only PR that fixes five sagemaker-train unit tests that silently issued real AWS calls (STS, IAM, SageMaker, S3). The fix moves @patch targets from the modules where names are defined to the modules where they are looked up (base_trainer, base_evaluator), adds a missing verify_evaluation_caller_permissions patch, and stubs the session in test_model_trainer_no_from_recipe_raises. I verified the import patterns in base_trainer.py (module-level from finetune_utils import ... and from ...defaults import TrainDefaults), the inside-function imports of _render_recipe_placeholders/_get_smtj_override_spec left correctly untouched, and base_evaluator.py's combined import of resolve_and_validate_role/verify_evaluation_caller_permissions. The changes are correct and align with established team guidance to patch at the import site in the module under test (Source: #6272). No issues found.

_{Reviewed commit 0a917e7. Automated review; verify before acting.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 72.89%. Comparing base (caf8a0e) to head (0a917e7).
⚠️ Report is 5 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6491      +/-   ##
==========================================
+ Coverage   64.26%   72.89%   +8.62%     
==========================================
  Files         570      570              
  Lines       71686    71686              
==========================================
+ Hits        46070    52252    +6182     
+ Misses      25616    19434    -6182     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.