Skip to content

Commit

Permalink
Clean-up
Browse files Browse the repository at this point in the history
  • Loading branch information
carl-baillargeon committed Jul 3, 2024
1 parent eaa6c75 commit a650846
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 9 deletions.
1 change: 1 addition & 0 deletions anta/catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -373,6 +373,7 @@ def from_list(data: ListAntaTestTuples) -> AntaCatalog:
raise
return AntaCatalog(tests)

# TODO: Move this change to a separate PR
@staticmethod
def merge(catalogs: list[AntaCatalog]) -> AntaCatalog:
"""Merge multiple AntaCatalog instances.
Expand Down
1 change: 1 addition & 0 deletions anta/device.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,7 @@ def __init__(
raise ValueError(message)
self.enable = enable
self._enable_password = enable_password
# TODO: Move the max_connections setting change to a separate PR
self._session: Device = Device(host=host, port=port, username=username, password=password, proto=proto, timeout=timeout, limits=Limits(max_connections=7))
ssh_params: dict[str, Any] = {}
if insecure:
Expand Down
17 changes: 9 additions & 8 deletions anta/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -592,12 +592,12 @@ async def wrapper(
self.save_commands_data(eos_data)
self.logger.debug("Test %s initialized with input data %s", self.name, eos_data)

# If the commands have not been collected, send them to the request manager and wait for the results
# If the commands have not been collected, send them to the test manager
if not self.collected:
logger.debug("<%s>: Sending commands for test %s to the Test Manager", self.device.name, self.name)
logger.debug("<%s>: Sending commands for test %s to the test manager", self.device.name, self.name)
condition = await self.send_commands()

# Wait until all commands have been collected
# Grab the condition returned from the manager and wait until all commands have been collected
async with condition:
await condition.wait_for(lambda: self.collected or self.result.result != "unset")
logger.debug("<%s>: Condition has been met for test %s", self.device.name, self.name)
Expand Down Expand Up @@ -697,11 +697,12 @@ async def get_commands(self) -> None:
commands = await asyncio.wait_for(get_await, timeout=2.0)
logger.debug("%d commands retrieved from the queue: %s", len(commands), commands)
condition = await self.parse_commands(commands)
# TODO: Put more info (context) in the condition
# TODO: Put more info (context) in the condition for logging
await self.notif_queue.put(condition)
except asyncio.TimeoutError: # noqa: PERF203
logger.warning("Timeout expired. Tests are done submitting commands.")
# Send the last batch
logger.debug("Timeout expired. Tests are done submitting commands.")

# Send the last batch if there are any commands left
if self.current_batch_commands:
logger.debug("Sending the last batch of commands")
await self.send_eapi_request(self.current_batch_id, self.current_batch_commands.copy())
Expand Down Expand Up @@ -731,7 +732,7 @@ async def parse_commands(self, commands: list[AntaCommand]) -> asyncio.Condition
return condition

async def send_eapi_request(self, batch_id: int, commands: list[AntaCommand]) -> None:
"""Send all the requests from the batches mapping."""
"""Send an eAPI request."""
eapi_request_id = f"Batch #{batch_id}"

logger.debug("Sending eAPI requests for batch %s with commands: %s", batch_id, commands)
Expand All @@ -743,7 +744,7 @@ async def send_eapi_request(self, batch_id: int, commands: list[AntaCommand]) ->

async def on_request_complete(self, batch_id: int) -> None:
"""TODO: Add docstring."""
# Notify the tests that the request is complete
# Notify the tests that the request is complete. Multiple tests can be waiting on the same batch (condition)
condition: asyncio.Condition = self.conditions[batch_id]
async with condition:
logger.debug("Notifying tests that the batch %s is complete. Condition: %s", batch_id, condition)
Expand Down
2 changes: 1 addition & 1 deletion anta/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,7 @@ def prepare_tests(
return device_to_tests

async def run_device_tests(device: AntaDevice, test_definitions: set[AntaTestDefinition], batch_size: int) -> list[TestResult]:
"""Run tests for a specific device using the RequestManager."""
"""Run tests for a specific device using the AntaTestManager."""
manager = AntaTestManager(device=device, batch_size=batch_size)
background_tasks = set()
coros = []
Expand Down

0 comments on commit a650846

Please sign in to comment.