|
9 | 9 | from faker import Faker |
10 | 10 | from fastapi import FastAPI, status |
11 | 11 | from fastapi.encoders import jsonable_encoder |
12 | | -from models_library.api_schemas_directorv2.dynamic_services import DynamicServiceGet |
| 12 | +from models_library.api_schemas_directorv2.dynamic_services import ( |
| 13 | + DynamicServiceGet, |
| 14 | + RetrieveDataOutEnveloped, |
| 15 | +) |
13 | 16 | from models_library.api_schemas_dynamic_scheduler.dynamic_services import ( |
14 | 17 | DynamicServiceStart, |
15 | 18 | DynamicServiceStop, |
@@ -55,14 +58,14 @@ def node_not_found(faker: Faker) -> NodeID: |
55 | 58 | @pytest.fixture |
56 | 59 | def service_status_new_style() -> DynamicServiceGet: |
57 | 60 | return TypeAdapter(DynamicServiceGet).validate_python( |
58 | | - DynamicServiceGet.model_config["json_schema_extra"]["examples"][1] |
| 61 | + DynamicServiceGet.model_json_schema()["examples"][1] |
59 | 62 | ) |
60 | 63 |
|
61 | 64 |
|
62 | 65 | @pytest.fixture |
63 | 66 | def service_status_legacy() -> NodeGet: |
64 | 67 | return TypeAdapter(NodeGet).validate_python( |
65 | | - NodeGet.model_config["json_schema_extra"]["examples"][1] |
| 68 | + NodeGet.model_json_schema()["examples"][1] |
66 | 69 | ) |
67 | 70 |
|
68 | 71 |
|
@@ -112,9 +115,7 @@ def mock_director_v2_service_state( |
112 | 115 | mock.get("/dynamic_services").respond( |
113 | 116 | status.HTTP_200_OK, |
114 | 117 | text=json.dumps( |
115 | | - jsonable_encoder( |
116 | | - DynamicServiceGet.model_config["json_schema_extra"]["examples"] |
117 | | - ) |
| 118 | + jsonable_encoder(DynamicServiceGet.model_json_schema()["examples"]) |
118 | 119 | ), |
119 | 120 | ) |
120 | 121 |
|
@@ -193,7 +194,7 @@ async def test_list_tracked_dynamic_services(rpc_client: RabbitMQRPCClient): |
193 | 194 | assert len(results) == 2 |
194 | 195 | assert results == [ |
195 | 196 | TypeAdapter(DynamicServiceGet).validate_python(x) |
196 | | - for x in DynamicServiceGet.model_config["json_schema_extra"]["examples"] |
| 197 | + for x in DynamicServiceGet.model_json_schema()["examples"] |
197 | 198 | ] |
198 | 199 |
|
199 | 200 |
|
@@ -223,7 +224,7 @@ async def test_get_state( |
223 | 224 | def dynamic_service_start() -> DynamicServiceStart: |
224 | 225 | # one for legacy and one for new style? |
225 | 226 | return TypeAdapter(DynamicServiceStart).validate_python( |
226 | | - DynamicServiceStart.model_config["json_schema_extra"]["example"] |
| 227 | + DynamicServiceStart.model_json_schema()["example"] |
227 | 228 | ) |
228 | 229 |
|
229 | 230 |
|
@@ -490,3 +491,38 @@ async def test_stop_dynamic_service_serializes_generic_errors( |
490 | 491 | ), |
491 | 492 | timeout_s=5, |
492 | 493 | ) |
| 494 | + |
| 495 | + |
| 496 | +@pytest.fixture |
| 497 | +def mock_director_v2_service_retrieve_inputs(node_id: NodeID) -> Iterator[None]: |
| 498 | + with respx.mock( |
| 499 | + base_url="http://director-v2:8000/v2", |
| 500 | + assert_all_called=False, |
| 501 | + assert_all_mocked=True, # IMPORTANT: KEEP always True! |
| 502 | + ) as mock: |
| 503 | + request_ok = mock.post(f"/dynamic_services/{node_id}:retrieve") |
| 504 | + |
| 505 | + request_ok.respond( |
| 506 | + status.HTTP_200_OK, |
| 507 | + text=TypeAdapter(RetrieveDataOutEnveloped) |
| 508 | + .validate_python( |
| 509 | + RetrieveDataOutEnveloped.model_json_schema()["examples"][0] |
| 510 | + ) |
| 511 | + .model_dump_json(), |
| 512 | + ) |
| 513 | + |
| 514 | + yield None |
| 515 | + |
| 516 | + |
| 517 | +async def test_retrieve_inputs( |
| 518 | + mock_director_v2_service_retrieve_inputs: None, |
| 519 | + rpc_client: RabbitMQRPCClient, |
| 520 | + node_id: NodeID, |
| 521 | +): |
| 522 | + results = await services.retrieve_inputs( |
| 523 | + rpc_client, node_id=node_id, port_keys=[], timeout_s=10 |
| 524 | + ) |
| 525 | + assert ( |
| 526 | + results.model_dump(mode="python") |
| 527 | + == RetrieveDataOutEnveloped.model_json_schema()["examples"][0] |
| 528 | + ) |
0 commit comments