Skip to content

Commit

Permalink
Fix ec2_client usage
Browse files Browse the repository at this point in the history
It is not callable.

Signed-off-by: Enrico Usai <[email protected]>
  • Loading branch information
enrico-usai committed Nov 7, 2023
1 parent 613da24 commit e235f4a
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 4 deletions.
4 changes: 2 additions & 2 deletions src/slurm_plugin/capacity_block_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -307,12 +307,12 @@ def _update_capacity_blocks_info_from_ec2(self, capacity_blocks: Dict[str, Capac
This method is called every time the CapacityBlockManager is re-initialized,
so when it starts/is restarted or when fleet configuration changes.
"""
capacity_block_ids = capacity_blocks.keys()
capacity_block_ids = list(capacity_blocks.keys())
logger.info("Retrieving Capacity Blocks information from EC2 for %s", ",".join(capacity_block_ids))
try:
capacity_block_reservations_info: List[
CapacityReservationInfo
] = self.ec2_client().describe_capacity_reservations(capacity_block_ids)
] = self.ec2_client.describe_capacity_reservations(capacity_block_ids)

for capacity_block_reservation_info in capacity_block_reservations_info:
capacity_block_id = capacity_block_reservation_info.capacity_reservation_id()
Expand Down
4 changes: 2 additions & 2 deletions tests/slurm_plugin/test_capacity_block_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ def test_get_reserved_nodenames(
)
mocked_client = mocker.MagicMock()
expected_exception = isinstance(capacity_blocks_info_from_ec2, Exception)
mocked_client.return_value.describe_capacity_reservations.side_effect = [
mocked_client.describe_capacity_reservations.side_effect = [
capacity_blocks_info_from_ec2 if expected_exception else capacity_blocks_info_from_ec2
]
capacity_block_manager._ec2_client = mocked_client
Expand Down Expand Up @@ -549,7 +549,7 @@ def test_update_capacity_blocks_info_from_ec2(
capacity_block_manager._capacity_blocks = init_capacity_blocks
expected_exception = isinstance(expected_error, AWSClientError)
mocked_client = mocker.MagicMock()
mocked_client.return_value.describe_capacity_reservations.side_effect = [
mocked_client.describe_capacity_reservations.side_effect = [
expected_error if expected_exception else capacity_blocks_info_from_ec2
]
capacity_block_manager._ec2_client = mocked_client
Expand Down

0 comments on commit e235f4a

Please sign in to comment.