def test_get_common_realtime_batch_params(api_byom_event, expected_common_realtime_batch_params):
realtime_event = api_byom_event("byom_realtime_custom", False)
batch_event = api_byom_event("byom_batch_custom", False)
realtime_event.update(batch_event)
TestCase().assertEqual(
get_common_realtime_batch_params(realtime_event, False, None), expected_common_realtime_batch_params
)
def test_validate(api_byom_event):
# event with required keys
valid_event = api_byom_event("byom_realtime_builtin")
TestCase().assertDictEqual(validate(valid_event), valid_event)
# event with missing required keys
bad_event = api_byom_event("byom_batch_builtin")
# remove required key
del bad_event["model_artifact_location"]
with pytest.raises(BadRequest) as execinfo:
validate(bad_event)
assert str(execinfo.value) == "Bad request. API body does not have the necessary parameter: model_artifact_location"
def test_get_image_uri(mocked_sm, api_byom_event):
custom_event = api_byom_event("byom_realtime_custom", False)
TestCase().assertEqual(get_image_uri("byom_realtime_custom", custom_event, "us-east-1"), "custom-image-uri")
mocked_sm.return_value = "test-imge-uri"
builtin_event = api_byom_event("byom_realtime_builtin", False)
TestCase().assertEqual(get_image_uri("byom_realtime_builtin", builtin_event, "us-east-1"), "test-imge-uri")
mocked_sm.assert_called_with(
framework=builtin_event.get("model_framework"),
region="us-east-1",
version=builtin_event.get("model_framework_version"),
)
def test_get_required_keys(
api_byom_event,
api_model_monitor_event,
required_api_byom_realtime_builtin,
required_api_byom_batch_builtin,
required_api_byom_realtime_custom,
required_api_byom_batch_custom,
required_api_keys_model_monitor,
):
# Required keys in byom, realtime, builtin
returned_keys = get_required_keys(api_byom_event("realtime", "xgboost"))
expected_keys = required_api_byom_realtime_builtin
TestCase().assertCountEqual(expected_keys, returned_keys)
# Required keys in byom, batch, builtin
returned_keys = get_required_keys(api_byom_event("batch", "xgboost"))
expected_keys = required_api_byom_batch_builtin
TestCase().assertCountEqual(expected_keys, returned_keys)
# Required keys in byom, realtime, custom
returned_keys = get_required_keys(api_byom_event("realtime", ""))
expected_keys = required_api_byom_realtime_custom
TestCase().assertCountEqual(expected_keys, returned_keys)
# Required keys in byom, batch, custom
returned_keys = get_required_keys(api_byom_event("batch", ""))
expected_keys = required_api_byom_batch_custom
TestCase().assertCountEqual(expected_keys, returned_keys)
# Required keys in model_monitor, default (no monitoring_type provided)
returned_keys = get_required_keys(api_model_monitor_event())
expected_keys = required_api_keys_model_monitor()
TestCase().assertCountEqual(expected_keys, returned_keys)
# Required keys in model_monitor, with monitoring_type provided
returned_keys = get_required_keys(api_model_monitor_event("modelquality"))
expected_keys = required_api_keys_model_monitor(False)
TestCase().assertCountEqual(expected_keys, returned_keys)
# assert for exceptions
with pytest.raises(BadRequest) as exceinfo:
get_required_keys({"pipeline_type": "not_supported"})
assert (
str(exceinfo.value) ==
"Bad request format. Pipeline type not supported. Check documentation for API & config formats"
)
with pytest.raises(BadRequest) as exceinfo:
get_required_keys({
"pipeline_type": "model_monitor",
"monitoring_type": "not_supported"
})
assert (
str(exceinfo.value) ==
"Bad request. MonitoringType supported are 'DataQuality'|'ModelQuality'|'ModelBias'|'ModelExplainability'"
)
with pytest.raises(BadRequest) as exceinfo:
get_required_keys({"pipeline_type": "byom"})
assert str(exceinfo.value) == "Bad request. missing keys for byom"
def test_provision_pipeline(api_image_builder_event, api_byom_event):
client = botocore.session.get_session().create_client("cloudformation")
stubber = Stubber(client)
expected_response = {
"statusCode": 200,
"isBase64Encoded": False,
"body": json.dumps({"message": "success: stack creation started", "pipeline_id": "1234"}),
"headers": {"Content-Type": content_type},
}
# The stubber will be called twice
stubber.add_response("create_stack", {"StackId": "1234"})
stubber.add_response("create_stack", {"StackId": "1234"})
with stubber:
with mock_s3():
event = api_image_builder_event
response = provision_pipeline(event, client)
assert response == expected_response
event = api_byom_event("byom_realtime_builtin")
s3_client = boto3.client("s3", region_name="us-east-1")
testfile = tempfile.NamedTemporaryFile()
s3_client.create_bucket(Bucket="testbucket")
upload_file_to_s3(
testfile.name, "testbucket", "blueprints/byom/byom_realtime_inference_pipeline.yaml", s3_client
)
s3_client.create_bucket(Bucket="testassetsbucket")
response = provision_pipeline(event, client, s3_client)
assert response == expected_response
def test_provision_pipeline(cf_client_params, api_byom_event):
client = botocore.session.get_session().create_client("cloudformation")
cp_client = botocore.session.get_session().create_client("codepipeline")
stubber = Stubber(client)
cp_stubber = Stubber(cp_client)
cfn_response = {"StackId": "1234"}
expected_params = cf_client_params
expected_response = {
"statusCode":
200,
"isBase64Encoded":
False,
"body":
json.dumps({
"message": "success: stack creation started",
"pipeline_id": "1234"
}),
"headers": {
"Content-Type": "plain/text"
},
}
event = api_byom_event("realtime", "xgboost")
cfn_response = {"StackId": "1234"}
stubber.add_response("create_stack", cfn_response, expected_params)
with stubber:
with cp_stubber:
response = provision_pipeline(event, client)
assert response == expected_response
def test_get_template_parameters(
api_byom_event,
api_image_builder_event,
expected_params_realtime_custom,
expected_image_builder_params,
expected_batch_params,
):
single_event = api_byom_event("byom_realtime_custom", False)
TestCase().assertEqual(get_template_parameters(single_event, False), expected_params_realtime_custom)
TestCase().assertEqual(get_template_parameters(api_image_builder_event, False), expected_image_builder_params)
TestCase().assertEqual(
get_template_parameters(api_byom_event("byom_batch_custom", False), False),
expected_batch_params,
)
# test for exception
with pytest.raises(BadRequest):
get_template_parameters({"pipeline_type": "unsupported"}, False)
def test_get_stack_name(api_byom_event, api_monitor_event, api_image_builder_event):
# realtime builtin pipeline
realtime_builtin = api_byom_event("byom_realtime_builtin")
assert (
get_stack_name(realtime_builtin)
== f"mlops-pipeline-{realtime_builtin['model_name']}-byompipelinerealtimebuiltin"
)
# batch builtin pipeline
batch_builtin = api_byom_event("byom_batch_builtin")
assert get_stack_name(batch_builtin) == f"mlops-pipeline-{batch_builtin['model_name']}-byompipelinebatchbuiltin"
# model monitor pipeline
assert get_stack_name(api_monitor_event) == f"mlops-pipeline-{api_monitor_event['model_name']}-byommodelmonitor"
# image builder pipeline
assert (
get_stack_name(api_image_builder_event)
== f"mlops-pipeline-{api_image_builder_event['image_tag']}-byompipelineimagebuilder"
)
def test_get_template_parameters(
template_parameters_realtime_builtin,
template_parameters_batch_builtin,
template_parameters_realtime_custom,
template_parameters_batch_custom,
template_parameters_model_monitor,
api_byom_event,
api_model_monitor_event,
get_parameters_keys,
):
# assert template parameters: realtime/builtin
_, returned_parameters = get_template_parameters(
api_byom_event("realtime", "xgboost"))
expected_parameters = template_parameters_realtime_builtin(
api_byom_event("realtime", "xgboost"))
TestCase().assertCountEqual(get_parameters_keys(expected_parameters),
get_parameters_keys(returned_parameters))
# assert template parameters: batch/builtin
_, returned_parameters = get_template_parameters(
api_byom_event("batch", "xgboost"))
expected_parameters = template_parameters_batch_builtin(
api_byom_event("batch", "xgboost"))
TestCase().assertCountEqual(get_parameters_keys(expected_parameters),
get_parameters_keys(returned_parameters))
# assert template parameters: realtime/custom
_, returned_parameters = get_template_parameters(
api_byom_event("realtime", ""))
expected_parameters = template_parameters_realtime_custom(
api_byom_event("realtime", ""))
TestCase().assertCountEqual(get_parameters_keys(expected_parameters),
get_parameters_keys(returned_parameters))
# assert template parameters: batch/custom
_, returned_parameters = get_template_parameters(
api_byom_event("batch", ""))
expected_parameters = template_parameters_batch_custom(
api_byom_event("batch", ""))
TestCase().assertCountEqual(get_parameters_keys(expected_parameters),
get_parameters_keys(returned_parameters))
# assert template parameters: model monitor
_, returned_parameters = get_template_parameters(api_model_monitor_event())
expected_parameters = template_parameters_model_monitor(
api_model_monitor_event())
TestCase().assertCountEqual(get_parameters_keys(expected_parameters),
get_parameters_keys(returned_parameters))
def test_get_bacth_specific_params(api_byom_event, expected_batch_specific_params):
batch_event = api_byom_event("byom_batch_custom", False)
TestCase().assertEqual(get_bacth_specific_params(batch_event, None), expected_batch_specific_params)
def test_get_realtime_specific_params(api_byom_event, expected_realtime_specific_params):
realtime_event = api_byom_event("byom_realtime_builtin", False)
TestCase().assertEqual(get_realtime_specific_params(realtime_event, None), expected_realtime_specific_params)
def test_get_stage_param(api_byom_event):
single_event = api_byom_event("byom_realtime_custom", False)
TestCase().assertEqual(get_stage_param(single_event, "data_capture_location", None), "bucket/datacapture")
multi_event = api_byom_event("byom_realtime_custom", True)
TestCase().assertEqual(get_stage_param(multi_event, "data_capture_location", "dev"), "bucket/dev_datacapture")
