def test_drain_machine_without_options(self):
self.setup_config_mock()
self.htcondor_adapter = HTCondorAdapter()
run_async(self.htcondor_adapter.drain_machine, drone_uuid='test')
self.mock_async_run_command.assert_called_with(
'condor_drain -graceful test')
def test_get_resource_ratios(self):
self.assertCountEqual(
list(
run_async(self.htcondor_adapter.get_resource_ratios,
drone_uuid="test")),
[self.cpu_ratio, self.memory_ratio],
)
self.mock_executor.return_value.run_command.assert_called_with(
self.command)
self.mock_executor.reset_mock()
self.assertEqual(
run_async(self.htcondor_adapter.get_resource_ratios,
drone_uuid="not_exists"),
[],
)
self.mock_executor.return_value.run_command.assert_not_called()
self.mock_executor.reset_mock()
self.assertEqual(
run_async(self.htcondor_adapter.get_resource_ratios,
drone_uuid="test_undefined"),
[],
)
self.mock_executor.return_value.run_command.assert_not_called()
self.mock_executor.reset_mock()
self.assertEqual(
run_async(self.htcondor_adapter.get_resource_ratios,
drone_uuid="test_error"),
[],
)
def test_run_command(self):
executor = SSHExecutor(host="test_host",
username="******",
client_keys=["TestKey"])
self.assertIsNone(
run_async(executor.run_command, command="Test").stdout)
self.mock_asyncssh.connect.assert_called_with(host="test_host",
username="******",
client_keys=["TestKey"])
self.mock_asyncssh.reset_mock()
executor = SSHExecutor(host="test_host",
username="******",
client_keys=("TestKey", ))
self.assertIsNone(
run_async(executor.run_command, command="Test").stdout)
self.mock_asyncssh.connect.assert_called_with(
host="test_host", username="******", client_keys=("TestKey", ))
self.mock_asyncssh.reset_mock()
executor = SSHExecutor(host="test_host",
username="******",
client_keys=("TestKey", ))
self.assertEqual(
run_async(executor.run_command, command="Test",
stdin_input="Test").stdout, "Test")
self.mock_asyncssh.connect.assert_called_with(
host="test_host", username="******", client_keys=("TestKey", ))
def test_terminate_resource(self):
run_async(
self.cloudstack_adapter.terminate_resource,
resource_attributes=AttributeDict(remote_resource_uuid='123456'))
self.mock_cloudstack_api.return_value.destroyVirtualMachine.assert_called_with(
id='123456')
def test_stop_resource(self):
run_async(
self.cloudstack_adapter.stop_resource,
resource_attributes=AttributeDict(remote_resource_uuid="123456"),
)
self.mock_cloudstack_api.return_value.stopVirtualMachine.assert_called_with(
id="123456")
def test_drain_machine_without_options(self):
self.setup_config_mock()
self.slurm_adapter = SlurmAdapter()
run_async(self.slurm_adapter.drain_machine, drone_uuid="VM-1")
self.mock_async_run_command.assert_called_with(
"scontrol update NodeName=host-10-18-1-1 State=DRAIN Reason='COBalD/TARDIS'"
)
def test_get_utilisation(self):
self.assertEqual(
run_async(self.fake_adapter.get_utilisation, "test-123"), 1.0)
self.config.BatchSystem.utilisation = AttributeDict(
get_value=lambda: 0.9)
self.fake_adapter = FakeBatchSystemAdapter()
self.assertEqual(
run_async(self.fake_adapter.get_utilisation, "test-123"), 0.9)
def test_drain_machine(self):
run_async(self.slurm_adapter.drain_machine, drone_uuid="VM-1")
self.mock_executor.return_value.run_command.assert_called_with(
"scontrol update NodeName=host-10-18-1-1 State=DRAIN Reason='COBalD/TARDIS'"
)
self.mock_executor.reset_mock()
self.assertIsNone(
run_async(self.slurm_adapter.drain_machine, drone_uuid="not_exists")
)
def test_get_utilisation(self):
self.assertEqual(
run_async(self.slurm_adapter.get_utilisation, drone_uuid="VM-1"),
min([self.cpu_ratio, self.memory_ratio]),
)
self.mock_executor.return_value.run_command.assert_called_with(self.command)
self.assertEqual(
run_async(self.slurm_adapter.get_utilisation, drone_uuid="not_exists"),
0.0,
)
def test_get_resource_state(self):
async def mocked_async_execute(sql_query: str, bind_parameters: dict):
return_values = {
"test-available-01234567ab": [{
"drone_uuid": "test-01234567ab",
"state": "AvailableState"
}],
"test-noexists-01234567ab": [],
}
return return_values[bind_parameters["drone_uuid"]]
self.sql_registry_mock.async_execute.side_effect = mocked_async_execute
self.assertEqual(
[],
run_async(
crud.get_resource_state,
sql_registry=self.sql_registry_mock,
drone_uuid="test-noexists-01234567ab",
),
)
self.sql_registry_mock.async_execute.assert_called_with(
"""
SELECT R.drone_uuid, RS.state
FROM Resources R
JOIN ResourceStates RS ON R.state_id = RS.state_id
WHERE R.drone_uuid = :drone_uuid""",
{"drone_uuid": "test-noexists-01234567ab"},
)
self.sql_registry_mock.reset_mock()
self.assertEqual(
[{
"drone_uuid": "test-01234567ab",
"state": "AvailableState"
}],
run_async(
crud.get_resource_state,
sql_registry=self.sql_registry_mock,
drone_uuid="test-available-01234567ab",
),
)
self.sql_registry_mock.async_execute.assert_called_with(
"""
SELECT R.drone_uuid, RS.state
FROM Resources R
JOIN ResourceStates RS ON R.state_id = RS.state_id
WHERE R.drone_uuid = :drone_uuid""",
{"drone_uuid": "test-available-01234567ab"},
)
def test_stop_resource(self):
run_async(
self.openstack_adapter.stop_resource,
resource_attributes=AttributeDict(
remote_resource_uuid="029312-1231-123123"),
)
params = {"os-stop": None}
self.mock_openstack_api.return_value.init_api.assert_called_with(
timeout=60)
self.mock_openstack_api.return_value.servers.run_action.assert_called_with(
"029312-1231-123123", **params)
def test_get_resource_ratios(self):
self.assertCountEqual(
list(
run_async(self.htcondor_adapter.get_resource_ratios,
drone_uuid='test')),
[self.cpu_ratio, self.memory_ratio])
self.mock_async_run_command.assert_called_with(self.command)
self.assertEqual(
run_async(self.htcondor_adapter.get_resource_ratios,
drone_uuid='not_exists'), {})
def test_terminate_resource(self):
run_async(
self.openstack_adapter.terminate_resource,
resource_attributes=AttributeDict(
remote_resource_uuid="029312-1231-123123"),
)
self.mock_openstack_api.return_value.init_api.assert_called_with(
timeout=60)
self.mock_openstack_api.return_value.servers.force_delete.assert_called_with(
"029312-1231-123123")
def test_get_machine_status(self):
self.assertEqual(
run_async(self.fake_adapter.get_machine_status, "test-123"),
MachineStatus.Available,
)
run_async(self.fake_adapter.drain_machine, "test-123")
self.assertEqual(
run_async(self.fake_adapter.get_machine_status, "test-123"),
MachineStatus.Drained,
)
def test_max_sessions(self):
with self.subTest(sessions="default"):
self.assertEqual(
DEFAULT_MAX_SESSIONS, run_async(probe_max_session, MockConnection())
)
for expected in (1, 9, 11, 20, 100):
with self.subTest(sessions=expected):
self.assertEqual(
expected,
run_async(probe_max_session, MockConnection(None, expected)),
)
def test_get_allocation(self):
self.assertEqual(
run_async(self.slurm_adapter.get_allocation, drone_uuid="VM-1"),
max([self.cpu_ratio, self.memory_ratio]),
)
self.mock_async_run_command.assert_called_with(self.command)
self.assertEqual(
run_async(self.slurm_adapter.get_allocation,
drone_uuid="not_exists"),
0.0,
)
def test_get_resource_ratios(self):
self.assertEqual(
list(run_async(self.slurm_adapter.get_resource_ratios, drone_uuid="VM-1")),
[self.cpu_ratio, self.memory_ratio],
)
self.mock_executor.return_value.run_command.assert_called_with(self.command)
self.mock_executor.reset_mock()
self.assertEqual(
run_async(self.slurm_adapter.get_resource_ratios, drone_uuid="not_exists"),
{},
)
def test_connection_property(self):
async def force_connection():
async with self.executor.bounded_connection as connection:
return connection
self.assertIsNone(self.executor._ssh_connection)
run_async(force_connection)
self.assertIsInstance(self.executor._ssh_connection, MockConnection)
current_ssh_connection = self.executor._ssh_connection
run_async(force_connection)
# make sure the connection is not needlessly replaced
self.assertEqual(self.executor._ssh_connection, current_ssh_connection)
def test_get_resource_state(self):
self.clear_lru_cache()
self.mock_crud.get_resource_state.return_value = async_return(
return_value=[{
"drone_uuid": "test-0123456789",
"state": "AvailableState"
}])
response = run_async(self.client.get,
"/resources/test-0123456789/state",
headers=self.headers)
self.assertEqual(response.status_code, 200)
self.assertEqual(
response.json(),
{
"drone_uuid": "test-0123456789",
"state": "AvailableState"
},
)
self.mock_crud.get_resource_state.return_value = async_return(
return_value=[])
response = run_async(self.client.get,
"/resources/test-1234567890/state",
headers=self.headers)
self.assertEqual(response.status_code, 404)
self.assertEqual(response.json(), {"detail": "Drone not found"})
response = run_async(self.client.get,
"/resources/test-invalid/state",
headers=self.headers)
self.assertEqual(response.status_code, 422)
self.assertEqual(
response.json(),
{
"detail": [{
"ctx": {
"pattern": "^\\S+-[A-Fa-f0-9]{10}$"
},
"loc": ["path", "drone_uuid"],
"msg":
'string does not match regex "^\\S+-[A-Fa-f0-9]{10}$"',
"type": "value_error.str.regex",
}]
},
)
response = run_async(self.client.get,
"/resources/state",
headers=self.headers)
self.assertEqual(response.status_code, 404)
self.assertEqual(response.json(), {"detail": "Not Found"})
def test_drain_machine_without_options(self):
self.setup_config_mock()
self.htcondor_adapter = HTCondorAdapter()
run_async(self.htcondor_adapter.drain_machine, drone_uuid="test")
self.mock_executor.return_value.run_command.assert_called_with(
"condor_drain -graceful slot1@test")
self.mock_executor.reset_mock()
run_async(self.htcondor_adapter.drain_machine, drone_uuid="test_uuid")
self.mock_executor.return_value.run_command.assert_called_with(
"condor_drain -graceful slot1@test_uuid@test")
def test_connection_property(self):
async def helper_coroutine():
return await self.executor.ssh_connection
self.assertIsNone(self.executor._ssh_connection)
run_async(helper_coroutine)
self.assertIsInstance(self.executor._ssh_connection, MockConnection)
current_ssh_connection = self.executor._ssh_connection
run_async(helper_coroutine)
self.assertEqual(self.executor._ssh_connection, current_ssh_connection)
def test_run_command(self):
self.assertEqual(
run_async(self.executor.run_command, 'exit 0').exit_code, 0)
with self.assertRaises(CommandExecutionFailure) as cf:
run_async(self.executor.run_command, 'exit 255')
self.assertEqual(cf.exception.exit_code, 255)
self.assertEqual(
run_async(self.executor.run_command, 'echo "Test"').stdout, "Test")
self.assertEqual(
run_async(self.executor.run_command,
'echo "Test" >>/dev/stderr').stderr, "Test")
def test_deploy_resource(self):
expected_resource_attributes = self.resource_attributes
expected_resource_attributes.update(created=datetime.now(),
updated=datetime.now())
return_resource_attributes = run_async(
self.moab_adapter.deploy_resource,
resource_attributes=AttributeDict(machine_type="test2large",
site_name="TestSite"),
)
if return_resource_attributes.created - expected_resource_attributes.created > timedelta(
seconds=1
) or return_resource_attributes.updated - expected_resource_attributes.updated > timedelta(
seconds=1):
raise Exception("Creation time or update time wrong!")
del (
expected_resource_attributes.created,
expected_resource_attributes.updated,
return_resource_attributes.created,
return_resource_attributes.updated,
)
self.assertEqual(return_resource_attributes,
expected_resource_attributes)
self.mock_executor.return_value.run_command.assert_called_with(
"msub -j oe -m p -l walltime=02:00:00:00,mem=120gb,nodes=1:ppn=20 startVM.py"
)
def test_get_allocation(self):
self.assertEqual(
run_async(self.htcondor_adapter.get_allocation, drone_uuid="test"),
max([self.cpu_ratio, self.memory_ratio]),
)
self.mock_executor.return_value.run_command.assert_called_with(
self.command)
def test_resource_status_raise(self):
# Update interval is 10 minutes, so set last update back by 2 minutes in order to execute sacct command and
# creation date to current date
created_timestamp = datetime.now()
new_timestamp = datetime.now() - timedelta(minutes=2)
self.moab_adapter._moab_status._last_update = new_timestamp
with self.assertRaises(TardisResourceStatusUpdateFailed):
run_async(
self.moab_adapter.resource_status,
AttributeDict(
resource_id=1351043,
remote_resource_uuid=1351043,
resource_state=ResourceStatus.Booting,
created=created_timestamp,
),
)
def test_get_utilisation(self):
self.assertEqual(
run_async(self.htcondor_adapter.get_utilisation,
drone_uuid="test"),
min([self.cpu_ratio, self.memory_ratio]),
)
self.mock_async_run_command.assert_called_with(self.command)
def test_resource_status(self):
self.assertEqual(run_async(self.openstack_adapter.resource_status,
resource_attributes=AttributeDict(remote_resource_uuid='029312-1231-123123')),
AttributeDict(drone_uuid='testsite-089123', remote_resource_uuid='029312-1231-123123',
resource_status=ResourceStatus.Running))
self.mock_openstack_api.return_value.init_api.assert_called_with(timeout=60)
self.mock_openstack_api.return_value.servers.get.assert_called_with('029312-1231-123123')
def test_get_resources(self):
self.clear_lru_cache()
full_expected_resources = [
{
"remote_resource_uuid": "14fa5640a7c146e482e8be41ec5dffea",
"state": "AvailableState",
"drone_uuid": "test-0125bc9fd8",
"site_name": "Test",
"machine_type": "m1.test",
"created": "2021-10-08T12:42:16.354400",
"updated": "2021-10-08T12:42:28.382025",
},
{
"remote_resource_uuid": "b3efcc5bc8b741af9222987e0434ca61",
"state": "AvailableState",
"drone_uuid": "test-6af3cfef14",
"site_name": "Test",
"machine_type": "m1.test",
"created": "2021-10-08T12:42:16.373454",
"updated": "2021-10-08T12:42:30.648325",
},
]
self.mock_crud.get_resources.return_value = async_return(
return_value=full_expected_resources)
response = run_async(self.client.get,
"/resources/",
headers=self.headers)
self.assertEqual(response.status_code, 200)
self.assertEqual(
response.json(),
full_expected_resources,
)
def test_resource_status_update(self):
self.assertEqual(self.resource_attributes["resource_status"],
ResourceStatus.Booting)
return_resource_attributes = run_async(
self.moab_adapter.resource_status,
resource_attributes=self.resource_attributes)
self.assertEqual(return_resource_attributes["resource_status"],
ResourceStatus.Running)
def test_construction_by_yaml(self):
executor = yaml.safe_load("""
!ShellExecutor
""")
self.assertEqual(
run_async(executor.run_command, "exit 0").exit_code, 0)
with self.assertRaises(CommandExecutionFailure) as cf:
run_async(self.executor.run_command, "exit 255")
self.assertEqual(cf.exception.exit_code, 255)
self.assertEqual(
run_async(executor.run_command, 'echo "Test"').stdout, "Test")
self.assertEqual(
run_async(executor.run_command,
'echo "Test" >>/dev/stderr').stderr, "Test")
