Skip to content

Commit a5f48e1

Browse files
llucaxllucax
committed
Simplify code using a safety timeout
Use `async with asyncio.timeout()` instead of using `asyncio.wait()` with extra tasks. Signed-off-by: Leandro Lucarella <[email protected]>
1 parent 7250a5d commit a5f48e1

File tree

1 file changed

+39
-149
lines changed

1 file changed

+39
-149
lines changed

‎tests/microgrid/power_distributing/test_power_distributing.py‎

Lines changed: 39 additions & 149 deletions
Original file line numberDiff line numberDiff line change
@@ -215,15 +215,9 @@ async def test_power_distributor_one_user(self, mocker: MockerFixture) -> None:
215215
await requests_channel.new_sender().send(request)
216216
result_rx = results_channel.new_receiver()
217217

218-
done, pending = await asyncio.wait(
219-
[asyncio.create_task(result_rx.receive())],
220-
timeout=SAFETY_TIMEOUT.total_seconds(),
221-
)
222-
223-
assert len(pending) == 0
224-
assert len(done) == 1
218+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
219+
result = await result_rx.receive()
225220

226-
result: Result = done.pop().result()
227221
assert isinstance(result, Success)
228222
assert result.succeeded_power.isclose(Power.from_kilowatts(1.0))
229223
assert result.excess_power.isclose(Power.from_watts(200.0))
@@ -284,15 +278,9 @@ async def test_power_distributor_exclusion_bounds(
284278
await requests_channel.new_sender().send(request)
285279
result_rx = results_channel.new_receiver()
286280

287-
done, pending = await asyncio.wait(
288-
[asyncio.create_task(result_rx.receive())],
289-
timeout=SAFETY_TIMEOUT.total_seconds(),
290-
)
281+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
282+
result = await result_rx.receive()
291283

292-
assert len(pending) == 0
293-
assert len(done) == 1
294-
295-
result: Result = done.pop().result()
296284
assert isinstance(result, Success)
297285
assert result.succeeded_power.isclose(Power.zero(), abs_tol=1e-9)
298286
assert result.excess_power.isclose(Power.zero(), abs_tol=1e-9)
@@ -308,15 +296,9 @@ async def test_power_distributor_exclusion_bounds(
308296
await requests_channel.new_sender().send(request)
309297
result_rx = results_channel.new_receiver()
310298

311-
done, pending = await asyncio.wait(
312-
[asyncio.create_task(result_rx.receive())],
313-
timeout=SAFETY_TIMEOUT.total_seconds(),
314-
)
315-
316-
assert len(pending) == 0
317-
assert len(done) == 1
299+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
300+
result = await result_rx.receive()
318301

319-
result = done.pop().result()
320302
assert isinstance(
321303
result, OutOfBounds
322304
), f"Expected OutOfBounds, got {result}"
@@ -380,15 +362,9 @@ async def test_two_batteries_one_inverters(self, mocker: MockerFixture) -> None:
380362
await requests_channel.new_sender().send(request)
381363
result_rx = results_channel.new_receiver()
382364

383-
done, pending = await asyncio.wait(
384-
[asyncio.create_task(result_rx.receive())],
385-
timeout=SAFETY_TIMEOUT.total_seconds(),
386-
)
387-
388-
assert len(pending) == 0
389-
assert len(done) == 1
365+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
366+
result = await result_rx.receive()
390367

391-
result: Result = done.pop().result()
392368
assert isinstance(result, Success)
393369
# Inverter bounded at 500
394370
assert result.succeeded_power.isclose(Power.from_watts(500.0))
@@ -459,15 +435,8 @@ async def test_two_batteries_one_broken_one_inverters(
459435
await requests_channel.new_sender().send(request)
460436
result_rx = results_channel.new_receiver()
461437

462-
done, pending = await asyncio.wait(
463-
[asyncio.create_task(result_rx.receive())],
464-
timeout=SAFETY_TIMEOUT.total_seconds(),
465-
)
466-
467-
assert len(pending) == 0
468-
assert len(done) == 1
469-
470-
result: Result = done.pop().result()
438+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
439+
result = await result_rx.receive()
471440

472441
assert isinstance(result, Error)
473442
assert result.request == request
@@ -516,15 +485,9 @@ async def test_battery_two_inverters(self, mocker: MockerFixture) -> None:
516485
await requests_channel.new_sender().send(request)
517486
result_rx = results_channel.new_receiver()
518487

519-
done, pending = await asyncio.wait(
520-
[asyncio.create_task(result_rx.receive())],
521-
timeout=SAFETY_TIMEOUT.total_seconds(),
522-
)
523-
524-
assert len(pending) == 0
525-
assert len(done) == 1
488+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
489+
result = await result_rx.receive()
526490

527-
result: Result = done.pop().result()
528491
assert isinstance(result, Success)
529492
# Inverters each bounded at 500, together 1000
530493
assert result.succeeded_power.isclose(Power.from_watts(1000.0))
@@ -569,15 +532,9 @@ async def test_two_batteries_three_inverters(self, mocker: MockerFixture) -> Non
569532
await requests_channel.new_sender().send(request)
570533
result_rx = results_channel.new_receiver()
571534

572-
done, pending = await asyncio.wait(
573-
[asyncio.create_task(result_rx.receive())],
574-
timeout=SAFETY_TIMEOUT.total_seconds(),
575-
)
576-
577-
assert len(pending) == 0
578-
assert len(done) == 1
535+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
536+
result = await result_rx.receive()
579537

580-
result: Result = done.pop().result()
581538
assert isinstance(result, Success)
582539
# each inverter is bounded at 500 and we have 3 inverters
583540
assert result.succeeded_power.isclose(Power.from_watts(1500.0))
@@ -656,15 +613,9 @@ async def test_two_batteries_one_inverter_different_exclusion_bounds_2(
656613
await requests_channel.new_sender().send(request)
657614
result_rx = results_channel.new_receiver()
658615

659-
done, pending = await asyncio.wait(
660-
[asyncio.create_task(result_rx.receive())],
661-
timeout=SAFETY_TIMEOUT.total_seconds(),
662-
)
663-
664-
assert len(pending) == 0
665-
assert len(done) == 1
616+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
617+
result = await result_rx.receive()
666618

667-
result: Result = done.pop().result()
668619
assert isinstance(result, OutOfBounds)
669620
assert result.request == request
670621
assert result.bounds == PowerBounds(-1000, -500, 500, 1000)
@@ -744,15 +695,9 @@ async def test_two_batteries_one_inverter_different_exclusion_bounds(
744695
await requests_channel.new_sender().send(request)
745696
result_rx = results_channel.new_receiver()
746697

747-
done, pending = await asyncio.wait(
748-
[asyncio.create_task(result_rx.receive())],
749-
timeout=SAFETY_TIMEOUT.total_seconds(),
750-
)
751-
752-
assert len(pending) == 0
753-
assert len(done) == 1
698+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
699+
result = await result_rx.receive()
754700

755-
result: Result = done.pop().result()
756701
assert isinstance(result, OutOfBounds)
757702
assert result.request == request
758703
# each inverter is bounded at 500
@@ -809,15 +754,9 @@ async def test_connected_but_not_requested_batteries(
809754
await requests_channel.new_sender().send(request)
810755
result_rx = results_channel.new_receiver()
811756

812-
done, pending = await asyncio.wait(
813-
[asyncio.create_task(result_rx.receive())],
814-
timeout=SAFETY_TIMEOUT.total_seconds(),
815-
)
816-
817-
assert len(pending) == 0
818-
assert len(done) == 1
757+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
758+
result = await result_rx.receive()
819759

820-
result: Result = done.pop().result()
821760
assert isinstance(result, Error)
822761
assert result.request == request
823762
assert (
@@ -865,15 +804,9 @@ async def test_battery_soc_nan(self, mocker: MockerFixture) -> None:
865804
await requests_channel.new_sender().send(request)
866805
result_rx = results_channel.new_receiver()
867806

868-
done, pending = await asyncio.wait(
869-
[asyncio.create_task(result_rx.receive())],
870-
timeout=SAFETY_TIMEOUT.total_seconds(),
871-
)
807+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
808+
result = await result_rx.receive()
872809

873-
assert len(pending) == 0
874-
assert len(done) == 1
875-
876-
result: Result = done.pop().result()
877810
assert isinstance(result, Success)
878811
assert result.succeeded_components == {19}
879812
assert result.succeeded_power.isclose(Power.from_watts(500.0))
@@ -921,15 +854,9 @@ async def test_battery_capacity_nan(self, mocker: MockerFixture) -> None:
921854
await requests_channel.new_sender().send(request)
922855
result_rx = results_channel.new_receiver()
923856

924-
done, pending = await asyncio.wait(
925-
[asyncio.create_task(result_rx.receive())],
926-
timeout=SAFETY_TIMEOUT.total_seconds(),
927-
)
857+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
858+
result = await result_rx.receive()
928859

929-
assert len(pending) == 0
930-
assert len(done) == 1
931-
932-
result: Result = done.pop().result()
933860
assert isinstance(result, Success)
934861
assert result.succeeded_components == {19}
935862
assert result.succeeded_power.isclose(Power.from_watts(500.0))
@@ -996,15 +923,9 @@ async def test_battery_power_bounds_nan(self, mocker: MockerFixture) -> None:
996923
await requests_channel.new_sender().send(request)
997924
result_rx = results_channel.new_receiver()
998925

999-
done, pending = await asyncio.wait(
1000-
[asyncio.create_task(result_rx.receive())],
1001-
timeout=SAFETY_TIMEOUT.total_seconds(),
1002-
)
1003-
1004-
assert len(pending) == 0
1005-
assert len(done) == 1
926+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
927+
result = await result_rx.receive()
1006928

1007-
result: Result = done.pop().result()
1008929
assert isinstance(result, Success)
1009930
assert result.succeeded_components == {19}
1010931
assert result.succeeded_power.isclose(Power.from_kilowatts(1.0))
@@ -1041,13 +962,9 @@ async def test_power_distributor_invalid_battery_id(
1041962
await requests_channel.new_sender().send(request)
1042963
result_rx = results_channel.new_receiver()
1043964

1044-
done, _ = await asyncio.wait(
1045-
[asyncio.create_task(result_rx.receive())],
1046-
timeout=SAFETY_TIMEOUT.total_seconds(),
1047-
)
965+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
966+
result = await result_rx.receive()
1048967

1049-
assert len(done) == 1
1050-
result: Result = done.pop().result()
1051968
assert isinstance(result, Error)
1052969
assert result.request == request
1053970
assert result.msg == "No battery 100, available batteries: 9, 19, 29"
@@ -1086,15 +1003,9 @@ async def test_power_distributor_one_user_adjust_power_consume(
10861003
await requests_channel.new_sender().send(request)
10871004
result_rx = results_channel.new_receiver()
10881005

1089-
done, pending = await asyncio.wait(
1090-
[asyncio.create_task(result_rx.receive())],
1091-
timeout=SAFETY_TIMEOUT.total_seconds(),
1092-
)
1093-
1094-
assert len(pending) == 0
1095-
assert len(done) == 1
1006+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
1007+
result = await result_rx.receive()
10961008

1097-
result = done.pop().result()
10981009
assert isinstance(result, OutOfBounds)
10991010
assert result is not None
11001011
assert result.request == request
@@ -1134,15 +1045,9 @@ async def test_power_distributor_one_user_adjust_power_supply(
11341045
await requests_channel.new_sender().send(request)
11351046
result_rx = results_channel.new_receiver()
11361047

1137-
done, pending = await asyncio.wait(
1138-
[asyncio.create_task(result_rx.receive())],
1139-
timeout=SAFETY_TIMEOUT.total_seconds(),
1140-
)
1048+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
1049+
result = await result_rx.receive()
11411050

1142-
assert len(pending) == 0
1143-
assert len(done) == 1
1144-
1145-
result = done.pop().result()
11461051
assert isinstance(result, OutOfBounds)
11471052
assert result is not None
11481053
assert result.request == request
@@ -1182,15 +1087,9 @@ async def test_power_distributor_one_user_adjust_power_success(
11821087
await requests_channel.new_sender().send(request)
11831088
result_rx = results_channel.new_receiver()
11841089

1185-
done, pending = await asyncio.wait(
1186-
[asyncio.create_task(result_rx.receive())],
1187-
timeout=SAFETY_TIMEOUT.total_seconds(),
1188-
)
1090+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
1091+
result = await result_rx.receive()
11891092

1190-
assert len(pending) == 0
1191-
assert len(done) == 1
1192-
1193-
result = done.pop().result()
11941093
assert isinstance(result, Success)
11951094
assert result.succeeded_power.isclose(Power.from_kilowatts(1.0))
11961095
assert result.excess_power.isclose(Power.zero(), abs_tol=1e-9)
@@ -1229,14 +1128,9 @@ async def test_not_all_batteries_are_working(self, mocker: MockerFixture) -> Non
12291128
await requests_channel.new_sender().send(request)
12301129
result_rx = results_channel.new_receiver()
12311130

1232-
done, pending = await asyncio.wait(
1233-
[asyncio.create_task(result_rx.receive())],
1234-
timeout=SAFETY_TIMEOUT.total_seconds(),
1235-
)
1131+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
1132+
result = await result_rx.receive()
12361133

1237-
assert len(pending) == 0
1238-
assert len(done) == 1
1239-
result = done.pop().result()
12401134
assert isinstance(result, Success)
12411135
assert result.succeeded_components == {19}
12421136
assert result.excess_power.isclose(Power.from_watts(700.0))
@@ -1284,13 +1178,9 @@ async def test_partial_failure_result(self, mocker: MockerFixture) -> None:
12841178
await requests_channel.new_sender().send(request)
12851179
result_rx = results_channel.new_receiver()
12861180

1287-
done, pending = await asyncio.wait(
1288-
[asyncio.create_task(result_rx.receive())],
1289-
timeout=SAFETY_TIMEOUT.total_seconds(),
1290-
)
1291-
assert len(pending) == 0
1292-
assert len(done) == 1
1293-
result = done.pop().result()
1181+
async with asyncio.timeout(SAFETY_TIMEOUT.total_seconds()):
1182+
result = await result_rx.receive()
1183+
12941184
assert isinstance(result, PartialFailure)
12951185
assert result.succeeded_components == batteries - failed_batteries
12961186
assert result.failed_components == failed_batteries

0 commit comments

Comments
 (0)