|
2 | 2 |
|
3 | 3 | import pytest |
4 | 4 |
|
5 | | -from inference_sdk.config import INTERNAL_REMOTE_EXEC_REQ_HEADER, execution_id |
| 5 | +from inference_sdk.config import ( |
| 6 | + EXECUTION_ID_HEADER, |
| 7 | + INTERNAL_REMOTE_EXEC_REQ_HEADER, |
| 8 | + execution_id, |
| 9 | +) |
6 | 10 | from inference_sdk.http.utils.request_building import ( |
7 | 11 | ImagePlacement, |
8 | 12 | RequestData, |
@@ -156,6 +160,43 @@ def test_prepare_requests_data() -> None: |
156 | 160 | ) |
157 | 161 |
|
158 | 162 |
|
| 163 | +def test_assembly_request_data_when_execution_id_is_set_and_headers_are_empty() -> None: |
| 164 | + # given |
| 165 | + token = execution_id.set("test-exec-id") |
| 166 | + |
| 167 | + try: |
| 168 | + # when |
| 169 | + result = assembly_request_data( |
| 170 | + url="https://some.com", |
| 171 | + batch_inference_inputs=[("image_1", None)], |
| 172 | + headers=None, |
| 173 | + parameters=None, |
| 174 | + payload=None, |
| 175 | + image_placement=ImagePlacement.DATA, |
| 176 | + ) |
| 177 | + |
| 178 | + # then |
| 179 | + assert result.headers == {EXECUTION_ID_HEADER: "test-exec-id"} |
| 180 | + finally: |
| 181 | + execution_id.reset(token) |
| 182 | + |
| 183 | + |
| 184 | +def test_assembly_request_data_when_image_placement_is_not_supported() -> None: |
| 185 | + # when / then |
| 186 | + with pytest.raises( |
| 187 | + NotImplementedError, |
| 188 | + match="Not implemented request building method", |
| 189 | + ): |
| 190 | + assembly_request_data( |
| 191 | + url="https://some.com", |
| 192 | + batch_inference_inputs=[("image_1", None)], |
| 193 | + headers=None, |
| 194 | + parameters=None, |
| 195 | + payload=None, |
| 196 | + image_placement=object(), |
| 197 | + ) |
| 198 | + |
| 199 | + |
159 | 200 | class TestInternalRemoteExecHeader: |
160 | 201 | @patch( |
161 | 202 | "inference_sdk.http.utils.request_building.ENABLE_INTERNAL_REMOTE_EXEC_HEADER", |
|
0 commit comments