tests

Author: kukushking

modules/sagemaker/sagemaker-model-monitoring/tests/test_baselining_handler.py

@@ -0,0 +1,123 @@
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def mock_sagemaker_modules():
+    """Mock sagemaker modules for these tests only."""
+    # Set environment variable
+    os.environ["SAGEMAKER_ROLE_ARN"] = "arn:aws:iam::123456789012:role/test-role"
+    
+    # Store original modules
+    original_modules = {}
+    sagemaker_module_names = [
+        "sagemaker",
+        "sagemaker.clarify",
+        "sagemaker.model_monitor",
+        "sagemaker.model_monitor.dataset_format",
+    ]
+    
+    for module_name in sagemaker_module_names:
+        if module_name in sys.modules:
+            original_modules[module_name] = sys.modules[module_name]
+    
+    # Mock the sagemaker modules
+    sys.modules["sagemaker"] = MagicMock()
+    sys.modules["sagemaker.clarify"] = MagicMock()
+    sys.modules["sagemaker.model_monitor"] = MagicMock()
+    sys.modules["sagemaker.model_monitor.dataset_format"] = MagicMock()
+    
+    # Add lambda path
+    lambda_path = "sagemaker_model_monitoring/lambda"
+    if lambda_path not in sys.path:
+        sys.path.insert(0, lambda_path)
+    
+    # Import after mocking
+    from baselining_handler import check_baselining_job, lambda_handler
+    
+    yield check_baselining_job, lambda_handler
+    
+    # Restore original modules
+    for module_name in sagemaker_module_names:
+        if module_name in original_modules:
+            sys.modules[module_name] = original_modules[module_name]
+        else:
+            sys.modules.pop(module_name, None)
+    
+    # Remove lambda path
+    if lambda_path in sys.path:
+        sys.path.remove(lambda_path)
+    
+    # Clean up imported handler module
+    if "baselining_handler" in sys.modules:
+        del sys.modules["baselining_handler"]
+
+
+@pytest.fixture
+def mock_env():
+    """Environment variables are already set in mock_sagemaker_modules."""
+    pass
+
+
+@pytest.fixture
+def mock_sagemaker_client():
+    """Mock boto3 SageMaker client."""
+    with patch("baselining_handler.sagemaker") as mock_client:
+        yield mock_client
+
+
+def test_lambda_handler_check_action(mock_sagemaker_modules, mock_env, mock_sagemaker_client):
+    """Test lambda_handler with check action."""
+    check_baselining_job, lambda_handler = mock_sagemaker_modules
+    event = {"action": "check", "job_name": "test-job"}
+
+    mock_sagemaker_client.describe_processing_job.return_value = {"ProcessingJobStatus": "Completed"}
+
+    result = lambda_handler(event, None)
+
+    assert result["statusCode"] == 200
+    assert result["status"] == "COMPLETED"
+    assert result["job_name"] == "test-job"
+
+
+def test_check_baselining_job_completed(mock_sagemaker_modules, mock_env, mock_sagemaker_client):
+    """Test check_baselining_job with completed job."""
+    check_baselining_job, lambda_handler = mock_sagemaker_modules
+    event = {"job_name": "test-job"}
+
+    mock_sagemaker_client.describe_processing_job.return_value = {"ProcessingJobStatus": "Completed"}
+
+    result = check_baselining_job(event)
+
+    assert result["statusCode"] == 200
+    assert result["status"] == "COMPLETED"
+    assert result["job_name"] == "test-job"
+
+
+def test_check_baselining_job_in_progress(mock_sagemaker_modules, mock_env, mock_sagemaker_client):
+    """Test check_baselining_job with in-progress job."""
+    check_baselining_job, lambda_handler = mock_sagemaker_modules
+    event = {"job_name": "test-job"}
+
+    mock_sagemaker_client.describe_processing_job.return_value = {"ProcessingJobStatus": "InProgress"}
+
+    result = check_baselining_job(event)
+
+    assert result["statusCode"] == 200
+    assert result["status"] == "IN_PROGRESS"
+    assert result["job_name"] == "test-job"
+
+
+def test_check_baselining_job_missing_job_name(mock_sagemaker_modules, mock_env):
+    """Test check_baselining_job with missing job_name."""
+    check_baselining_job, lambda_handler = mock_sagemaker_modules
+    event = {}
+
+    result = check_baselining_job(event)
+
+    assert result["statusCode"] == 400
+    assert "error" in result
+    assert "Missing required parameter" in result["error"]

modules/sagemaker/sagemaker-model-monitoring/tests/test_stack.py

@@ -22,6 +22,8 @@ def stack_model_package_input(
     enable_model_quality_monitor: bool = False,
     enable_model_bias_monitor: bool = False,
     enable_model_explainability_monitor: bool = False,
+    baseline_training_data_s3_uri: str = None,
+    baseline_output_data_s3_uri: str = None,
 ) -> cdk.Stack:
     from sagemaker_model_monitoring import settings, stack
 
@@ -51,6 +53,8 @@ def stack_model_package_input(
         enable_model_quality_monitor=enable_model_quality_monitor,
         enable_model_bias_monitor=enable_model_bias_monitor,
         enable_model_explainability_monitor=enable_model_explainability_monitor,
+        baseline_training_data_s3_uri=baseline_training_data_s3_uri,
+        baseline_output_data_s3_uri=baseline_output_data_s3_uri,
     )
 
     return stack.SageMakerModelMonitoringStack(
@@ -88,6 +92,36 @@ def test_synthesize_stack_model_explainability(stack_defaults: None) -> None:
     template.resource_count_is("AWS::SageMaker::ModelExplainabilityJobDefinition", 1)
 
 
+def test_baseline_generation_resources(stack_defaults: None) -> None:
+    """Test that baseline generation resources are created when URIs are provided."""
+    stack = stack_model_package_input(
+        enable_data_quality_monitor=True,
+        baseline_training_data_s3_uri="s3://test-bucket/training-data.csv",
+        baseline_output_data_s3_uri="s3://test-bucket/baselines/",
+    )
+    template = Template.from_stack(stack)
+
+    # Check for Lambda function
+    template.resource_count_is("AWS::Lambda::Function", 1)
+
+    # Check for Step Functions state machine
+    template.resource_count_is("AWS::StepFunctions::StateMachine", 1)
+
+    # Check for EventBridge rule
+    template.resource_count_is("AWS::Events::Rule", 1)
+
+
+def test_no_baseline_generation_without_uris(stack_defaults: None) -> None:
+    """Test that baseline generation resources are NOT created without URIs."""
+    stack = stack_model_package_input(enable_data_quality_monitor=True)
+    template = Template.from_stack(stack)
+
+    # Should not have Lambda, Step Functions, or EventBridge resources
+    template.resource_count_is("AWS::Lambda::Function", 0)
+    template.resource_count_is("AWS::StepFunctions::StateMachine", 0)
+    template.resource_count_is("AWS::Events::Rule", 0)
+
+
 def test_no_cdk_nag_errors(stack_defaults: None) -> None:
     stack = stack_model_package_input(
         enable_data_quality_monitor=True,
@@ -102,3 +136,19 @@ def test_no_cdk_nag_errors(stack_defaults: None) -> None:
         Match.string_like_regexp(r"AwsSolutions-.*"),
     )
     assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors"
+
+
+def test_no_cdk_nag_errors_with_baseline(stack_defaults: None) -> None:
+    """Test CDK nag with baseline generation enabled."""
+    stack = stack_model_package_input(
+        enable_data_quality_monitor=True,
+        baseline_training_data_s3_uri="s3://test-bucket/training-data.csv",
+        baseline_output_data_s3_uri="s3://test-bucket/baselines/",
+    )
+    cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks())
+
+    nag_errors = Annotations.from_stack(stack).find_error(
+        "*",
+        Match.string_like_regexp(r"AwsSolutions-.*"),
+    )
+    assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors"