|
| 1 | +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"). You |
| 4 | +# may not use this file except in compliance with the License. A copy of |
| 5 | +# the License is located at |
| 6 | +# |
| 7 | +# http://aws.amazon.com/apache2.0/ |
| 8 | +# |
| 9 | +# or in the "license" file accompanying this file. This file is |
| 10 | +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF |
| 11 | +# ANY KIND, either express or implied. See the License for the specific |
| 12 | +# language governing permissions and limitations under the License. |
| 13 | +"""Unit tests for _resolve_data_cache_config and _resolve_container_spec.""" |
| 14 | +from __future__ import absolute_import |
| 15 | + |
| 16 | +import pytest |
| 17 | + |
| 18 | +from sagemaker.core.shapes import ( |
| 19 | + InferenceComponentDataCacheConfig, |
| 20 | + InferenceComponentContainerSpecification, |
| 21 | +) |
| 22 | +from sagemaker.serve.model_builder_utils import _ModelBuilderUtils |
| 23 | + |
| 24 | + |
| 25 | +class ConcreteUtils(_ModelBuilderUtils): |
| 26 | + """Concrete class to test mixin methods.""" |
| 27 | + pass |
| 28 | + |
| 29 | + |
| 30 | +@pytest.fixture |
| 31 | +def utils(): |
| 32 | + return ConcreteUtils() |
| 33 | + |
| 34 | + |
| 35 | +# ============================================================ |
| 36 | +# Tests for _resolve_data_cache_config |
| 37 | +# ============================================================ |
| 38 | + |
| 39 | +class TestResolveDataCacheConfig: |
| 40 | + def test_none_returns_none(self, utils): |
| 41 | + assert utils._resolve_data_cache_config(None) is None |
| 42 | + |
| 43 | + def test_already_typed_passthrough(self, utils): |
| 44 | + config = InferenceComponentDataCacheConfig(enable_caching=True) |
| 45 | + result = utils._resolve_data_cache_config(config) |
| 46 | + assert result is config |
| 47 | + assert result.enable_caching is True |
| 48 | + |
| 49 | + def test_dict_with_enable_caching_true(self, utils): |
| 50 | + result = utils._resolve_data_cache_config({"enable_caching": True}) |
| 51 | + assert isinstance(result, InferenceComponentDataCacheConfig) |
| 52 | + assert result.enable_caching is True |
| 53 | + |
| 54 | + def test_dict_with_enable_caching_false(self, utils): |
| 55 | + result = utils._resolve_data_cache_config({"enable_caching": False}) |
| 56 | + assert isinstance(result, InferenceComponentDataCacheConfig) |
| 57 | + assert result.enable_caching is False |
| 58 | + |
| 59 | + def test_dict_missing_enable_caching_raises(self, utils): |
| 60 | + with pytest.raises(ValueError, match="must contain the required 'enable_caching' key"): |
| 61 | + utils._resolve_data_cache_config({}) |
| 62 | + |
| 63 | + def test_dict_with_extra_keys_still_works(self, utils): |
| 64 | + """Extra keys are ignored; only enable_caching is required.""" |
| 65 | + result = utils._resolve_data_cache_config( |
| 66 | + {"enable_caching": True, "extra_key": "ignored"} |
| 67 | + ) |
| 68 | + assert isinstance(result, InferenceComponentDataCacheConfig) |
| 69 | + assert result.enable_caching is True |
| 70 | + |
| 71 | + def test_invalid_type_raises(self, utils): |
| 72 | + with pytest.raises(ValueError, match="data_cache_config must be a dict"): |
| 73 | + utils._resolve_data_cache_config("invalid") |
| 74 | + |
| 75 | + def test_invalid_type_int_raises(self, utils): |
| 76 | + with pytest.raises(ValueError, match="data_cache_config must be a dict"): |
| 77 | + utils._resolve_data_cache_config(42) |
| 78 | + |
| 79 | + def test_invalid_type_list_raises(self, utils): |
| 80 | + with pytest.raises(ValueError, match="data_cache_config must be a dict"): |
| 81 | + utils._resolve_data_cache_config([True]) |
| 82 | + |
| 83 | + |
| 84 | +# ============================================================ |
| 85 | +# Tests for _resolve_container_spec |
| 86 | +# ============================================================ |
| 87 | + |
| 88 | +class TestResolveContainerSpec: |
| 89 | + def test_none_returns_none(self, utils): |
| 90 | + assert utils._resolve_container_spec(None) is None |
| 91 | + |
| 92 | + def test_already_typed_passthrough(self, utils): |
| 93 | + spec = InferenceComponentContainerSpecification( |
| 94 | + image="my-image:latest", |
| 95 | + artifact_url="s3://bucket/artifact", |
| 96 | + environment={"KEY": "VALUE"}, |
| 97 | + ) |
| 98 | + result = utils._resolve_container_spec(spec) |
| 99 | + assert result is spec |
| 100 | + |
| 101 | + def test_dict_full(self, utils): |
| 102 | + result = utils._resolve_container_spec({ |
| 103 | + "image": "my-image:latest", |
| 104 | + "artifact_url": "s3://bucket/artifact", |
| 105 | + "environment": {"KEY": "VALUE"}, |
| 106 | + }) |
| 107 | + assert isinstance(result, InferenceComponentContainerSpecification) |
| 108 | + assert result.image == "my-image:latest" |
| 109 | + assert result.artifact_url == "s3://bucket/artifact" |
| 110 | + assert result.environment == {"KEY": "VALUE"} |
| 111 | + |
| 112 | + def test_dict_image_only(self, utils): |
| 113 | + result = utils._resolve_container_spec({"image": "my-image:latest"}) |
| 114 | + assert isinstance(result, InferenceComponentContainerSpecification) |
| 115 | + assert result.image == "my-image:latest" |
| 116 | + |
| 117 | + def test_dict_artifact_url_only(self, utils): |
| 118 | + result = utils._resolve_container_spec({"artifact_url": "s3://bucket/model.tar.gz"}) |
| 119 | + assert isinstance(result, InferenceComponentContainerSpecification) |
| 120 | + assert result.artifact_url == "s3://bucket/model.tar.gz" |
| 121 | + |
| 122 | + def test_dict_environment_only(self, utils): |
| 123 | + result = utils._resolve_container_spec({"environment": {"A": "B"}}) |
| 124 | + assert isinstance(result, InferenceComponentContainerSpecification) |
| 125 | + assert result.environment == {"A": "B"} |
| 126 | + |
| 127 | + def test_dict_empty(self, utils): |
| 128 | + """Empty dict creates a spec with no fields set.""" |
| 129 | + result = utils._resolve_container_spec({}) |
| 130 | + assert isinstance(result, InferenceComponentContainerSpecification) |
| 131 | + |
| 132 | + def test_dict_with_extra_keys(self, utils): |
| 133 | + """Extra keys are ignored.""" |
| 134 | + result = utils._resolve_container_spec({ |
| 135 | + "image": "img", |
| 136 | + "unknown_key": "ignored", |
| 137 | + }) |
| 138 | + assert isinstance(result, InferenceComponentContainerSpecification) |
| 139 | + assert result.image == "img" |
| 140 | + |
| 141 | + def test_invalid_type_raises(self, utils): |
| 142 | + with pytest.raises(ValueError, match="container must be a dict"): |
| 143 | + utils._resolve_container_spec("invalid") |
| 144 | + |
| 145 | + def test_invalid_type_int_raises(self, utils): |
| 146 | + with pytest.raises(ValueError, match="container must be a dict"): |
| 147 | + utils._resolve_container_spec(123) |
| 148 | + |
| 149 | + def test_invalid_type_list_raises(self, utils): |
| 150 | + with pytest.raises(ValueError, match="container must be a dict"): |
| 151 | + utils._resolve_container_spec([{"image": "img"}]) |
0 commit comments