async def _run_sub_test(wait_for_containers: bool) -> None:
ip = "192.0.2.0"
container = ContainerInstance(instance_id="123", ip_address=ip)
mock_onedocker_service.start_containers = MagicMock(
return_value=[container])
mock_onedocker_service.wait_for_pending_containers = AsyncMock(
return_value=[container])
container.status = (ContainerInstanceStatus.COMPLETED
if wait_for_containers else
ContainerInstanceStatus.STARTED)
mock_wait_for_containers_async.return_value = [container]
with patch.object(PIDProtocolRunStage,
"files_exist") as mock_files_exist, patch.object(
PIDProtocolRunStage,
"put_server_ips") as mock_put_server_ips:
mock_files_exist.return_value = True
num_shards = 2
input_path = "in"
output_path = "out"
# Run publisher
publisher_run_stage = PIDProtocolRunStage(
stage=UnionPIDStage.PUBLISHER_RUN_PID,
instance_repository=mock_instance_repo,
storage_svc=mock_storage_service,
onedocker_svc=mock_onedocker_service,
onedocker_binary_config=self.onedocker_binary_config,
)
instance_id = "123"
stage_input = PIDStageInput(
input_paths=[input_path],
output_paths=[output_path],
num_shards=num_shards,
instance_id=instance_id,
)
# if we are waiting for containers, then the stage should finish
# otherwise, it should start and then return
self.assertEqual(
PIDStageStatus.COMPLETED
if wait_for_containers else PIDStageStatus.STARTED,
await publisher_run_stage.run(
stage_input=stage_input,
wait_for_containers=wait_for_containers,
),
)
# Check create_instances_async was called with the correct parameters
if wait_for_containers:
mock_wait_for_containers_async.assert_called_once()
else:
mock_wait_for_containers_async.assert_not_called()
mock_onedocker_service.start_containers.assert_called_once()
(
_,
called_kwargs,
) = mock_onedocker_service.start_containers.call_args_list[0]
self.assertEqual(num_shards,
len(called_kwargs["cmd_args_list"]))
# Check `put_payload` was called with the correct parameters
mock_put_server_ips.assert_called_once_with(
instance_id=instance_id, server_ips=[ip])
# if wait for containers is False, there are 4 updates.
# if wait_for_containers is True, then there is another update
# that updates the instance status and containers to complete, so 5
mock_instance_repo.read.assert_called_with(instance_id)
self.assertEqual(mock_instance_repo.read.call_count,
4 + int(wait_for_containers))
self.assertEqual(
mock_instance_repo.update.call_count,
4 + int(wait_for_containers),
)
async def _run_sub_test(wait_for_containers: bool, ) -> None:
ip = "192.0.2.0"
container = ContainerInstance(instance_id="123", ip_address=ip)
mock_onedocker_svc.start_containers = MagicMock(
return_value=[container])
mock_onedocker_svc.wait_for_pending_containers = AsyncMock(
return_value=[container])
container.status = (ContainerInstanceStatus.COMPLETED
if wait_for_containers else
ContainerInstanceStatus.STARTED)
mock_wait_for_containers_async.return_value = [container]
test_onedocker_binary_config = OneDockerBinaryConfig(
tmp_directory="/test_tmp_directory/",
binary_version="latest",
repository_path="test_path/",
)
stage = PIDShardStage(
stage=UnionPIDStage.PUBLISHER_SHARD,
instance_repository=mock_instance_repo,
storage_svc=mock_storage_svc,
onedocker_svc=mock_onedocker_svc,
onedocker_binary_config=test_onedocker_binary_config,
)
instance_id = "444"
stage_input = PIDStageInput(
input_paths=["in"],
output_paths=["out"],
num_shards=123,
instance_id=instance_id,
)
# Basic test: All good
with patch.object(PIDShardStage, "files_exist") as mock_fe:
mock_fe.return_value = True
stage = PIDShardStage(
stage=UnionPIDStage.PUBLISHER_SHARD,
instance_repository=mock_instance_repo,
storage_svc=mock_storage_svc,
onedocker_svc=mock_onedocker_svc,
onedocker_binary_config=test_onedocker_binary_config,
)
status = await stage.run(
stage_input, wait_for_containers=wait_for_containers)
self.assertEqual(
PIDStageStatus.COMPLETED
if wait_for_containers else PIDStageStatus.STARTED,
status,
)
mock_onedocker_svc.start_containers.assert_called_once()
if wait_for_containers:
mock_wait_for_containers_async.assert_called_once()
else:
mock_wait_for_containers_async.assert_not_called()
# instance status is updated to READY, STARTED, then COMPLETED
mock_instance_repo.read.assert_called_with(instance_id)
self.assertEqual(mock_instance_repo.read.call_count, 4)
self.assertEqual(mock_instance_repo.update.call_count, 4)
# Input not ready
with patch.object(PIDShardStage, "files_exist") as mock_fe:
mock_fe.return_value = False
status = await stage.run(
stage_input, wait_for_containers=wait_for_containers)
self.assertEqual(PIDStageStatus.FAILED, status)
# Multiple input paths (invariant exception)
with patch.object(PIDShardStage, "files_exist") as mock_fe:
with self.assertRaises(ValueError):
mock_fe.return_value = True
stage_input.input_paths = ["in1", "in2"]
stage = PIDShardStage(
stage=UnionPIDStage.PUBLISHER_SHARD,
instance_repository=mock_instance_repo,
storage_svc=mock_storage_svc,
onedocker_svc=mock_onedocker_svc,
onedocker_binary_config=test_onedocker_binary_config,
)
status = await stage.run(
stage_input, wait_for_containers=wait_for_containers)
self.assertEqual(
PIDStageStatus.COMPLETED
if wait_for_containers else PIDStageStatus.STARTED,
status,
)
async def _run_sub_test(wait_for_containers: bool, ) -> None:
ip = "192.0.2.0"
container = ContainerInstance(
instance_id="123",
ip_address=ip,
status=ContainerInstanceStatus.STARTED)
mock_onedocker_svc.start_containers = MagicMock(
return_value=[container])
mock_onedocker_svc.wait_for_pending_containers = AsyncMock(
return_value=[container])
container.status = (ContainerInstanceStatus.COMPLETED
if wait_for_containers else
ContainerInstanceStatus.STARTED)
mock_wait_for_containers_async.return_value = [container]
stage = PIDPrepareStage(
stage=UnionPIDStage.PUBLISHER_PREPARE,
instance_repository=mock_instance_repo,
storage_svc=mock_storage_svc,
onedocker_svc=mock_onedocker_svc,
onedocker_binary_config=MagicMock(
task_definition="offline-task:1#container",
tmp_directory="/tmp/",
binary_version="latest",
),
)
instance_id = "444"
stage_input = PIDStageInput(
input_paths=["in"],
output_paths=["out"],
num_shards=2,
instance_id=instance_id,
)
# Basic test: All good
with patch.object(PIDPrepareStage, "files_exist") as mock_fe:
mock_fe.return_value = True
stage = PIDPrepareStage(
stage=UnionPIDStage.PUBLISHER_PREPARE,
instance_repository=mock_instance_repo,
storage_svc=mock_storage_svc,
onedocker_svc=mock_onedocker_svc,
onedocker_binary_config=MagicMock(
task_definition="offline-task:1#container",
tmp_directory="/tmp/",
binary_version="latest",
),
)
status = await stage.run(
stage_input, wait_for_containers=wait_for_containers)
self.assertEqual(
PIDStageStatus.COMPLETED
if wait_for_containers else PIDStageStatus.STARTED,
status,
)
self.assertEqual(
mock_onedocker_svc.start_containers.call_count, 2)
if wait_for_containers:
self.assertEqual(mock_wait_for_containers_async.call_count,
2)
else:
mock_wait_for_containers_async.assert_not_called()
mock_instance_repo.read.assert_called_with(instance_id)
self.assertEqual(mock_instance_repo.read.call_count, 4)
self.assertEqual(mock_instance_repo.update.call_count, 4)
with patch.object(PIDPrepareStage,
"files_exist") as mock_fe, patch.object(
PIDPrepareStage, "prepare") as mock_prepare:
mock_fe.return_value = True
status = await stage.run(
stage_input, wait_for_containers=wait_for_containers)
mock_prepare.assert_called_with(instance_id, "in", "out", 2,
wait_for_containers, None)
# Input not ready
with patch.object(PIDPrepareStage, "files_exist") as mock_fe:
mock_fe.return_value = False
status = await stage.run(
stage_input, wait_for_containers=wait_for_containers)
self.assertEqual(PIDStageStatus.FAILED, status)
# Multiple input paths (invariant exception)
with patch.object(PIDPrepareStage, "files_exist") as mock_fe:
with self.assertRaises(ValueError):
mock_fe.return_value = True
stage_input.input_paths = ["in1", "in2"]
stage = PIDPrepareStage(
stage=UnionPIDStage.PUBLISHER_PREPARE,
instance_repository=mock_instance_repo,
storage_svc=mock_storage_svc,
onedocker_svc=mock_onedocker_svc,
onedocker_binary_config=MagicMock(
task_definition="offline-task:1#container",
tmp_directory="/tmp/",
binary_version="latest",
),
)
status = await stage.run(
stage_input, wait_for_containers=wait_for_containers)