def test_terminate_with_concurrent_read(self):
scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16', 1, 'Succeeded')
updated_scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16', 0, 'Succeeded')
instance = AzureScaleSetInstance('fake_id', 'fake_vm', datetime.now())
future = TestingFuture()
mock_api = mock.Mock(AzureApi)
mock_api.list_scale_sets = mock.Mock(return_value=[scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[instance])
mock_api.terminate_scale_set_instances = mock.Mock(return_value=future)
cached_api = AzureWriteThroughCachedApi(mock_api)
self.assertEqual(cached_api.list_scale_sets('test_rg'), [scale_set])
self.assertEqual(cached_api.list_scale_set_instances(scale_set), [instance])
cached_api.terminate_scale_set_instances(scale_set, [instance])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(scale_set)
mock_api.terminate_scale_set_instances.assert_called_once_with(scale_set, [instance])
# Call list again concurrently with the delete, and make sure it's still served from the cache
self.assertEqual(cached_api.list_scale_sets('test_rg'), [scale_set])
self.assertEqual(cached_api.list_scale_set_instances(scale_set), [instance])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(scale_set)
future.complete()
mock_api.list_scale_sets = mock.Mock(return_value=[updated_scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[])
self.assertEqual(cached_api.list_scale_sets('test_rg'), [updated_scale_set])
self.assertEqual(cached_api.list_scale_set_instances(updated_scale_set), [])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(updated_scale_set)
def test_failed_scale_up(self):
region = 'test'
mock_client, monitor_client, resource_client = _default_mock_clients(
region)
instance_type = 'Standard_D1_v2'
resource_group = 'test-resource-group'
failed_scale_set = AzureScaleSet(region, resource_group,
'test-scale-set1', instance_type, 0,
'Failed')
scale_set = AzureScaleSet(region, resource_group, 'test-scale-set2',
instance_type, 0, 'Succeeded')
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group,
AzureWrapper(mock_client, monitor_client, resource_client),
instance_type, False, [failed_scale_set, scale_set], [])
virtual_scale_set.scale(5)
mock_client.virtual_machine_scale_sets.create_or_update.assert_called_once(
)
self.assertEqual(
mock_client.virtual_machine_scale_sets.create_or_update.
call_args[1]['parameters'].sku.capacity, 5)
def test_slow_scale_up(self):
region = 'test'
mock_client, monitor_client, resource_client = _default_mock_clients(
region)
instance_type = 'Standard_D1_v2'
resource_group = 'test-resource-group'
scale_set = AzureScaleSet(region, resource_group, 'test-scale-set',
instance_type, 0, 'Succeeded')
scale_set2 = AzureScaleSet(region, resource_group, 'test-scale-set2',
instance_type, 0, 'Succeeded')
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group,
AzureWrapper(mock_client, monitor_client, resource_client),
instance_type, True, [scale_set, scale_set2], [])
virtual_scale_set.scale(2)
self.assertEqual(
mock_client.virtual_machine_scale_sets.create_or_update.call_count,
2)
self.assertEqual(
mock_client.virtual_machine_scale_sets.create_or_update.
call_args_list[0][1]['parameters'].sku.capacity, 1)
self.assertEqual(
mock_client.virtual_machine_scale_sets.create_or_update.
call_args_list[1][1]['parameters'].sku.capacity, 1)
def test_update(self):
scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16', 1, 'Succeeded')
updated_scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16', 0, 'Succeeded')
instance = AzureScaleSetInstance('fake_id', 'fake_vm', datetime.now())
future = CompletedFuture(None)
mock_api = mock.Mock(AzureApi)
mock_api.list_scale_sets = mock.Mock(return_value=[scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[instance])
mock_api.update_scale_set = mock.Mock(return_value=future)
cached_api = AzureWriteThroughCachedApi(mock_api)
self.assertEqual(cached_api.list_scale_sets('test_rg'), [scale_set])
self.assertEqual(cached_api.list_scale_set_instances(scale_set), [instance])
cached_api.update_scale_set(scale_set, 0).result()
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(scale_set)
mock_api.update_scale_set.assert_called_once_with(scale_set, 0)
mock_api.list_scale_sets = mock.Mock(return_value=[updated_scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[])
self.assertEqual(cached_api.list_scale_sets('test_rg'), [updated_scale_set])
self.assertEqual(cached_api.list_scale_set_instances(updated_scale_set), [])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(updated_scale_set)
def test_inconsistent_delegate(self):
scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16', 0, 'Succeeded')
updated_scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16', 1, 'Succeeded')
instance = AzureScaleSetInstance('fake_id', 'fake_vm', datetime.now())
future = CompletedFuture(None)
mock_api = mock.Mock(AzureApi)
mock_api.list_scale_sets = mock.Mock(return_value=[scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[])
mock_api.update_scale_set = mock.Mock(return_value=future)
cached_api = AzureWriteThroughCachedApi(mock_api)
self.assertEqual(cached_api.list_scale_sets('test_rg'), [scale_set])
self.assertEqual(cached_api.list_scale_set_instances(scale_set), [])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(scale_set)
cached_api.update_scale_set(scale_set, 1).result()
mock_api.update_scale_set.assert_called_once_with(scale_set, 1)
mock_api.list_scale_sets = mock.Mock(return_value=[updated_scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[])
self.assertEqual(cached_api.list_scale_sets('test_rg'), [updated_scale_set])
self.assertEqual(cached_api.list_scale_set_instances(updated_scale_set), [])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(updated_scale_set)
# Test that even if there is inconsistency between the list_scale_sets and list_scale_set_instances, the
# cache doesn't end up with bad data
mock_api.list_scale_set_instances = mock.Mock(return_value=[instance])
self.assertEqual(cached_api.list_scale_set_instances(updated_scale_set), [instance])
mock_api.list_scale_set_instances.assert_called_once_with(updated_scale_set)
def test_refresh(self):
scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16',
1, 'Succeeded')
updated_scale_set = AzureScaleSet('eastus', 'test_rg', 'test',
'Standard_H16', 0, 'Succeeded')
scale_set2 = AzureScaleSet('eastus', 'test_rg', 'test2',
'Standard_H16', 0, 'Succeeded')
instance = AzureScaleSetInstance('fake_id', 'fake_vm', datetime.now())
mock_api = mock.Mock(AzureApi)
mock_api.list_scale_sets = mock.Mock(return_value=[scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[instance])
cached_api = AzureWriteThroughCachedApi(mock_api)
self.assertEqual(cached_api.list_scale_sets('test_rg'), [scale_set])
self.assertEqual(cached_api.list_scale_set_instances(scale_set),
[instance])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(scale_set)
mock_api.list_scale_sets = mock.Mock(
return_value=[updated_scale_set, scale_set2])
mock_api.list_scale_set_instances = mock.Mock(return_value=[])
self.assertEqual(
set(cached_api.list_scale_sets('test_rg', force_refresh=True)),
{updated_scale_set, scale_set2})
self.assertEqual(
cached_api.list_scale_set_instances(updated_scale_set), [])
mock_api.list_scale_sets.assert_called_once_with('test_rg')
mock_api.list_scale_set_instances.assert_called_once_with(
updated_scale_set)
def test_scale_in(self):
region = 'test'
resource_group = 'test-resource-group'
instance = VirtualMachineScaleSetVM(location=region)
instance.vm_id = 'test-vm-id'
instance.instance_id = 0
instance.instance_view = VirtualMachineInstanceView()
instance.instance_view.statuses = []
mock_client, monitor_client, resource_client = _default_mock_clients(
region, instances=[instance])
TestNode = collections.namedtuple('TestNode',
['instance_id', 'unschedulable'])
test_node = TestNode(instance_id=instance.vm_id, unschedulable=False)
instance_type = 'Standard_D1_v2'
scale_set = AzureScaleSet(region, resource_group, 'test-scale-set',
instance_type, 1, 'Succeeded')
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group,
AzureWrapper(mock_client, monitor_client, resource_client),
instance_type, False, [scale_set], [test_node])
self.assertEqual(virtual_scale_set.instance_ids, {instance.vm_id})
self.assertEqual(virtual_scale_set.nodes, [test_node])
virtual_scale_set.scale_nodes_in([test_node])
mock_client.virtual_machine_scale_sets.delete_instances.assert_called_once_with(
resource_group, scale_set.name, [instance.instance_id])
def test_tainted_scale_set(self):
region = 'test'
mock_client, monitor_client, resource_client = _default_mock_clients(
region)
instance_type = 'Standard_NC24'
resource_group = 'test-resource-group'
scale_set = AzureScaleSet(region,
resource_group,
'test-scale-set',
instance_type,
0,
'Succeeded',
no_schedule_taints={'gpu': 'yes'})
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group,
AzureWrapper(mock_client, monitor_client, resource_client),
instance_type, True, [scale_set], [])
dir_path = os.path.dirname(os.path.realpath(__file__))
with open(os.path.join(dir_path, 'data/busybox.yaml'), 'r') as f:
dummy_pod = yaml.load(f.read())
pod = KubePod(pykube.Pod(None, dummy_pod))
self.assertFalse(virtual_scale_set.is_taints_tolerated(pod))
dummy_pod['spec']['tolerations'] = [{
'key': 'gpu',
'operator': 'Exists'
}]
pod = KubePod(pykube.Pod(None, dummy_pod))
self.assertTrue(virtual_scale_set.is_taints_tolerated(pod))
def test_scale_up(self):
region = 'test'
mock_client = mock.Mock()
mock_client.virtual_machine_scale_set_vms = mock.Mock()
mock_client.virtual_machine_scale_set_vms.list = mock.Mock(
return_value=[])
mock_client.virtual_machine_scale_sets = mock.Mock()
mock_client.virtual_machine_scale_sets.create_or_update = mock.Mock()
monitor_client = mock.Mock()
monitor_client.activity_logs = mock.Mock()
monitor_client.activity_logs.list = mock.Mock(return_value=[])
instance_type = 'Standard_D1_v2'
resource_group = 'test-resource-group'
scale_set = AzureScaleSet(region, resource_group, 'test-scale-set',
instance_type, 0, 'Succeeded')
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group, AzureWrapper(mock_client, monitor_client),
instance_type, False, [scale_set], [])
virtual_scale_set.scale(5)
mock_client.virtual_machine_scale_sets.create_or_update.assert_called_once(
)
self.assertEqual(
mock_client.virtual_machine_scale_sets.create_or_update.
call_args[1]['parameters'].sku.capacity, 5)
def test_out_of_quota(self):
region = 'test'
mock_client = mock.Mock()
mock_client.virtual_machine_scale_set_vms = mock.Mock()
mock_client.virtual_machine_scale_set_vms.list = mock.Mock(
return_value=[])
mock_client.virtual_machine_scale_sets = mock.Mock()
mock_client.virtual_machine_scale_sets.create_or_update = mock.Mock()
monitor_client = mock.Mock()
monitor_client.activity_logs = mock.Mock()
monitor_client.activity_logs.list = mock.Mock(return_value=[])
instance_type = 'Standard_D1_v2'
resource_group = 'test-resource-group'
scale_set = AzureScaleSet(region,
resource_group,
'test-scale-set',
instance_type,
0,
'Succeeded',
timeout_until=datetime.now(pytz.utc) +
timedelta(minutes=10),
timeout_reason="fake reason")
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group, AzureWrapper(mock_client, monitor_client),
instance_type, False, [scale_set], [])
self.assertTrue(virtual_scale_set.is_timed_out())
def test_out_of_quota(self):
scale_set = VirtualMachineScaleSet('eastus', {}, sku=Sku('Standard_H16', capacity=1))
scale_set.name = 'test'
scale_set.provisioning_state = 'Succeeded'
scale_set.id = 'fake_id'
compute_client = mock.Mock()
compute_client.virtual_machine_scale_sets = mock.Mock()
compute_client.virtual_machine_scale_sets.list = mock.Mock(return_value=[scale_set])
reason = "Operation results in exceeding quota limits of Core. Maximum allowed: 800, Current in use: 784, Additional requested: 320."
message = "{\"error\":{\"code\":\"OperationNotAllowed\",\"message\":\"" + reason + "\"}}"
monitor_client = mock.Mock()
monitor_client.activity_logs = mock.Mock()
now = datetime.now(pytz.utc)
monitor_client.activity_logs.list = mock.Mock(return_value=[EventData('Error',
now,
now,
resource_id=scale_set.id,
status=LocalizableString('Failed'),
properties={'statusCode': 'Conflict',
'statusMessage': message})])
api = AzureWrapper(compute_client, monitor_client)
resource_group = 'test_rg'
expected = AzureScaleSet(scale_set.location, resource_group, scale_set.name, scale_set.sku.name, scale_set.sku.capacity,
scale_set.provisioning_state, now + TIMEOUT_PERIOD, reason)
acutal = api.list_scale_sets(resource_group)
self.assertEqual([expected], acutal)
compute_client.virtual_machine_scale_sets.list.assert_called_once_with(resource_group)
monitor_client.activity_logs.list.assert_called_once()
def test_basic(self):
scale_set = VirtualMachineScaleSet(
'eastus', {
PRIORITY_TAG: '1',
NO_SCHEDULE_TAINTS_TAG: json.dumps({'gpu': 'yes'})
},
sku=Sku('Standard_H16', capacity=1))
scale_set.name = 'test'
scale_set.provisioning_state = 'Succeeded'
scale_set.id = 'fake_id'
compute_client = mock.Mock()
compute_client.virtual_machine_scale_sets = mock.Mock()
compute_client.virtual_machine_scale_sets.list = mock.Mock(
return_value=[scale_set])
monitor_client = mock.Mock()
monitor_client.activity_logs = mock.Mock()
monitor_client.activity_logs.list = mock.Mock(return_value=[])
api = AzureWrapper(compute_client, monitor_client, None)
resource_group = 'test_rg'
expected = AzureScaleSet(scale_set.location,
resource_group,
scale_set.name,
scale_set.sku.name,
scale_set.sku.capacity,
scale_set.provisioning_state,
priority=1,
no_schedule_taints={'gpu': 'yes'})
self.assertEqual([expected], api.list_scale_sets(resource_group))
compute_client.virtual_machine_scale_sets.list.assert_called_once_with(
resource_group)
monitor_client.activity_logs.list.assert_called_once()
def test_priority(self):
region = 'test'
mock_client = mock.Mock()
mock_client.virtual_machine_scale_set_vms = mock.Mock()
mock_client.virtual_machine_scale_set_vms.list = mock.Mock(
return_value=[])
mock_client.virtual_machine_scale_sets = mock.Mock()
mock_client.virtual_machine_scale_sets.create_or_update = mock.Mock()
monitor_client = mock.Mock()
monitor_client.activity_logs = mock.Mock()
monitor_client.activity_logs.list = mock.Mock(return_value=[])
instance_type = 'Standard_D1_v2'
resource_group = 'test-resource-group'
scale_set = AzureScaleSet(region,
resource_group,
'test-scale-set',
instance_type,
0,
'Succeeded',
priority=-1)
# Name sorts lexicographically before previous scale set, but priority is after it
scale_set2 = AzureScaleSet(region,
resource_group,
'a-test-scale-set',
instance_type,
0,
'Succeeded',
priority=1)
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group, AzureWrapper(mock_client, monitor_client),
instance_type, True, [scale_set, scale_set2], [])
virtual_scale_set.scale(1)
self.assertEqual(virtual_scale_set.global_priority, -1)
self.assertEqual(
mock_client.virtual_machine_scale_sets.create_or_update.call_count,
1)
self.assertEqual(
mock_client.virtual_machine_scale_sets.create_or_update.
call_args_list[0][0][1], 'test-scale-set')
def test_copied(self):
scale_set = AzureScaleSet('eastus', 'test_rg', 'test', 'Standard_H16', 1, 'Succeeded')
instance = AzureScaleSetInstance('fake_id', 'fake_vm', datetime.now())
mock_api = mock.Mock(AzureApi)
mock_api.list_scale_sets = mock.Mock(return_value=[scale_set])
mock_api.list_scale_set_instances = mock.Mock(return_value=[instance])
cached_api = AzureWriteThroughCachedApi(mock_api)
returned_scale_set = cached_api.list_scale_sets('test_rg')[0]
self.assertEqual(returned_scale_set.capacity, 1)
returned_scale_set.capacity = 0
self.assertEqual(cached_api.list_scale_sets('test_rg')[0].capacity, 1)
returned_instance = cached_api.list_scale_set_instances(scale_set)[0]
self.assertEqual(returned_instance.vm_id, 'fake_vm')
returned_instance.vm_id = 'modified'
self.assertEqual(cached_api.list_scale_set_instances(scale_set)[0].vm_id, 'fake_vm')
def test_out_of_quota(self):
region = 'test'
mock_client, monitor_client, resource_client = _default_mock_clients(
region)
instance_type = 'Standard_D1_v2'
resource_group = 'test-resource-group'
scale_set = AzureScaleSet(region,
resource_group,
'test-scale-set',
instance_type,
0,
'Succeeded',
timeout_until=datetime.now(pytz.utc) +
timedelta(minutes=10),
timeout_reason="fake reason")
virtual_scale_set = AzureVirtualScaleSet(
region, resource_group,
AzureWrapper(mock_client, monitor_client, resource_client),
instance_type, False, [scale_set], [])
self.assertTrue(virtual_scale_set.is_timed_out())