253 строки
11 KiB
Python
253 строки
11 KiB
Python
# --------------------------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License. See License.txt in the project root for license information.
|
|
# --------------------------------------------------------------------------------------------
|
|
|
|
import sys
|
|
import os
|
|
import logging
|
|
import datetime
|
|
|
|
try:
|
|
import unittest2 as unittest
|
|
except ImportError:
|
|
import unittest
|
|
|
|
try:
|
|
from unittest import mock
|
|
except ImportError:
|
|
import mock
|
|
|
|
from ui_pools import PoolsUI, AzureBatchPoolInfo
|
|
from pools import AzureBatchPools
|
|
from environment import AzureBatchEnvironment
|
|
import azure.batch_extensions as batch
|
|
from azure.batch_extensions import models
|
|
|
|
|
|
class AzureTestBatchPools(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.mock_self = mock.create_autospec(AzureBatchPools)
|
|
self.mock_self._log = logging.getLogger("TestPools")
|
|
self.mock_self.batch = mock.create_autospec(batch.BatchExtensionsClient)
|
|
self.mock_self.batch.pool = mock.create_autospec(batch.operations.ExtendedPoolOperations)
|
|
self.mock_self.environment = mock.create_autospec(AzureBatchEnvironment)
|
|
return super(AzureTestBatchPools, self).setUp()
|
|
|
|
@mock.patch("pools.PoolsUI")
|
|
def test_pools_initialize(self, mock_ui):
|
|
pools = AzureBatchPools(4, "frame", "call")
|
|
mock_ui.assert_called_with(pools, "frame")
|
|
|
|
def test_pools_configure(self):
|
|
session = mock.Mock(batch="batch")
|
|
AzureBatchPools.configure(self.mock_self, session, "env_manager")
|
|
self.assertEqual("batch", self.mock_self.batch)
|
|
|
|
def test_pools_list_pools(self):
|
|
self.mock_self.batch = mock.create_autospec(batch.BatchExtensionsClient)
|
|
self.mock_self.batch.pool = mock.create_autospec(batch.operations.ExtendedPoolOperations)
|
|
pool1 = mock.create_autospec(models.CloudPool)
|
|
pool1.id = "12345"
|
|
pool1.virtual_machine_configuration = mock.create_autospec(models.VirtualMachineConfiguration)
|
|
pool1.virtual_machine_configuration.image_reference = mock.create_autospec(models.ImageReference)
|
|
pool1.virtual_machine_configuration.image_reference.publisher = "MicrosoftWindows"
|
|
pool1.creation_time = datetime.datetime.now()
|
|
pool2 = mock.create_autospec(models.CloudPool)
|
|
pool2.id = "67890"
|
|
pool2.virtual_machine_configuration = mock.create_autospec(models.VirtualMachineConfiguration)
|
|
pool2.virtual_machine_configuration.image_reference = mock.create_autospec(models.ImageReference)
|
|
pool2.virtual_machine_configuration.image_reference.publisher = "LinuxUbuntu"
|
|
pool2.creation_time = datetime.datetime.now()
|
|
self.mock_self._call = lambda x: [pool1, pool2]
|
|
|
|
ids = AzureBatchPools.list_pools(self.mock_self)
|
|
self.assertEqual(ids, ["67890", "12345"])
|
|
self.assertEqual(len(self.mock_self.pools), 2)
|
|
|
|
pool1.id = "Maya_Pool_A"
|
|
pool1.virtual_machine_configuration.image_reference.publisher = "batch"
|
|
pool2.id = "Maya_Auto_Pool_B"
|
|
pool2.virtual_machine_configuration.image_reference.publisher = "batch"
|
|
ids = AzureBatchPools.list_pools(self.mock_self)
|
|
self.assertEqual(ids, ["Maya_Pool_A"])
|
|
self.assertEqual(len(self.mock_self.pools), 2)
|
|
|
|
def test_pools_get_pools(self):
|
|
def list_pools():
|
|
self.mock_self.pools = [mock.Mock(id="1234", display_name="name")]
|
|
self.mock_self.list_pools = list_pools
|
|
self.mock_self.ui = mock.create_autospec(PoolsUI)
|
|
self.mock_self.ui.create_pool_entry.return_value = "pool_entry"
|
|
|
|
displayed = AzureBatchPools.get_pools(self.mock_self)
|
|
self.assertEqual(displayed, ["pool_entry"])
|
|
self.assertEqual(len(self.mock_self.pools), 1)
|
|
self.mock_self.ui.create_pool_entry.assert_called_with('name', 0)
|
|
|
|
def test_pools_pool_selected(self):
|
|
pool_ui = mock.create_autospec(AzureBatchPoolInfo)
|
|
pool_ui.index = 0
|
|
self.mock_self.selected_pool = None
|
|
|
|
AzureBatchPools.pool_selected(self.mock_self, None)
|
|
self.assertFalse(self.mock_self.update_pool.call_count)
|
|
AzureBatchPools.pool_selected(self.mock_self, pool_ui)
|
|
self.mock_self.update_pool.assert_called_with(0)
|
|
AzureBatchPools.pool_selected(self.mock_self, pool_ui)
|
|
pool_ui.collapse.assert_called_with()
|
|
|
|
|
|
@mock.patch("pools.maya")
|
|
def test_pools_update_pool(self, mock_maya):
|
|
self.mock_self.batch = mock.create_autospec(batch.BatchExtensionsClient)
|
|
self.mock_self.batch.pool = mock.Mock(get="get")
|
|
self.mock_self.batch.compute_node = mock.Mock(list="list")
|
|
pool = mock.create_autospec(models.CloudPool)
|
|
pool.application_licenses = ["maya", "arnold"]
|
|
pool.display_name = "name"
|
|
pool.current_dedicated_nodes = 3
|
|
pool.target_dedicated_nodes = 5
|
|
pool.state = mock.Mock(value="resizing")
|
|
pool.id = "Maya_Pool_12345"
|
|
pool.max_tasks_per_node = 1
|
|
pool.allocation_state = mock.Mock(value="allocating")
|
|
pool.creation_time = datetime.datetime.now()
|
|
pool.vm_size = "Standard_A1"
|
|
pool.virtual_machine_configuration = mock.create_autospec(batch.models.VirtualMachineConfiguration)
|
|
pool.virtual_machine_configuration.image_reference = "image"
|
|
pool.virtual_machine_configuration.container_configuration = None
|
|
self.mock_self.environment.get_image_label.return_value = "Batch Windows Image"
|
|
self.mock_self.pools = [pool]
|
|
pool_ui = mock.create_autospec(AzureBatchPoolInfo)
|
|
pool_ui.index = 0
|
|
|
|
def call(func, pool_id):
|
|
self.assertEqual(pool_id, 'Maya_Pool_12345')
|
|
if func == "get":
|
|
return pool
|
|
else:
|
|
return []
|
|
self.mock_self._call = call
|
|
self.mock_self.selected_pool = pool_ui
|
|
self.mock_self.ui = mock.create_autospec(PoolsUI)
|
|
|
|
AzureBatchPools.update_pool(self.mock_self, 0)
|
|
self.assertEqual(self.mock_self.ui.refresh.call_count, 0)
|
|
self.assertEqual(mock_maya.refresh.call_count, 2)
|
|
AzureBatchPools.update_pool(self.mock_self, 1)
|
|
self.assertEqual(self.mock_self.ui.refresh.call_count, 1)
|
|
mock_maya.refresh.side_effect = ValueError("Error")
|
|
AzureBatchPools.update_pool(self.mock_self, 0)
|
|
self.assertEqual(self.mock_self.ui.refresh.call_count, 2)
|
|
|
|
def test_pools_auto_pool(self):
|
|
self.mock_self.pools = [mock.Mock(id='Maya_Pool_123'), mock.Mock(id='Maya_Auto_Pool_123')]
|
|
self.mock_self.selected_pool = mock.Mock(index=1)
|
|
self.assertTrue(AzureBatchPools.is_auto_pool(self.mock_self))
|
|
self.mock_self.selected_pool = mock.Mock(index=0)
|
|
self.assertFalse(AzureBatchPools.is_auto_pool(self.mock_self))
|
|
self.mock_self.selected_pool = mock.Mock(index=2)
|
|
self.assertFalse(AzureBatchPools.is_auto_pool(self.mock_self))
|
|
|
|
def test_pools_get_size(self):
|
|
|
|
self.mock_self.pools = [mock.Mock(target_dedicated_nodes=5, target_low_priority_nodes=3),
|
|
mock.Mock(target_dedicated_nodes="8", target_low_priority_nodes="3"),
|
|
mock.Mock(target_dedicated_nodes=None, target_low_priority_nodes=None)]
|
|
|
|
self.mock_self.selected_pool = mock.Mock(index=0)
|
|
self.assertEqual(AzureBatchPools.get_pool_size(self.mock_self), (5, 3))
|
|
|
|
self.mock_self.selected_pool = mock.Mock(index=1)
|
|
self.assertEqual(AzureBatchPools.get_pool_size(self.mock_self), (8, 3))
|
|
|
|
self.mock_self.selected_pool = mock.Mock(index=2)
|
|
self.assertEqual(AzureBatchPools.get_pool_size(self.mock_self), (0, 0))
|
|
|
|
self.mock_self.selected_pool = mock.Mock(index=3)
|
|
self.assertEqual(AzureBatchPools.get_pool_size(self.mock_self), (0, 0))
|
|
|
|
def test_pools_delete(self):
|
|
|
|
mock_pool = mock.create_autospec(models.CloudPool)
|
|
mock_pool.id = "pool id"
|
|
self.mock_self.pools = [mock_pool]
|
|
|
|
def call(func, *args):
|
|
self.assertTrue(callable(func))
|
|
return func(*args)
|
|
|
|
self.mock_self._call = call
|
|
self.mock_self.selected_pool = mock.Mock(index=0)
|
|
self.mock_self.ui = mock.create_autospec(PoolsUI)
|
|
|
|
AzureBatchPools.delete_pool(self.mock_self)
|
|
self.assertEqual(self.mock_self.ui.refresh.call_count, 1)
|
|
self.mock_self.batch.pool.delete.assert_called_with("pool id")
|
|
|
|
self.mock_self.batch.pool.delete.side_effect = AttributeError("boom")
|
|
AzureBatchPools.delete_pool(self.mock_self)
|
|
self.assertEqual(self.mock_self.ui.refresh.call_count, 2)
|
|
self.mock_self.batch.pool.delete.assert_called_with("pool id")
|
|
|
|
self.mock_self.selected_pool.index = 1
|
|
AzureBatchPools.delete_pool(self.mock_self)
|
|
self.assertEqual(self.mock_self.ui.refresh.call_count, 3)
|
|
self.assertEqual(self.mock_self.batch.pool.delete.call_count, 2)
|
|
|
|
def test_pools_create(self):
|
|
pool_obj = None
|
|
def call(func, new_pool):
|
|
global pool_obj
|
|
self.assertTrue(callable(func))
|
|
pool_obj = new_pool
|
|
self.assertEqual(new_pool.target_dedicated_nodes, 3)
|
|
self.assertEqual(new_pool.target_low_priority_nodes, 5)
|
|
self.assertEqual(new_pool.display_name, "Maya Pool for test job")
|
|
self.assertEqual(new_pool.application_licenses, ['maya'])
|
|
self.assertEqual(new_pool.virtual_machine_configuration.node_agent_sku_id, 'sku_id')
|
|
self.assertEqual(new_pool.virtual_machine_configuration.image_reference.publisher, 'foo')
|
|
return func(new_pool)
|
|
|
|
self.mock_self._call = call
|
|
self.mock_self.environment.get_application_licenses.return_value = ['maya']
|
|
|
|
image_ref = mock.create_autospec(models.ImageReference)
|
|
image_ref.publisher = 'foo'
|
|
vm_config = mock.create_autospec(models.VirtualMachineConfiguration)
|
|
vm_config.image_reference = image_ref
|
|
vm_config.node_agent_sku_id = 'sku_id'
|
|
|
|
self.mock_self.environment.build_virtualmachineconfiguration.return_value = vm_config
|
|
|
|
AzureBatchPools.create_pool(self.mock_self, (3, 5), "test job")
|
|
self.mock_self.batch.pool.add.assert_called_with(mock.ANY)
|
|
|
|
@mock.patch("pools.maya")
|
|
def test_pools_resize(self, mock_maya):
|
|
|
|
def call(func, *args):
|
|
self.assertTrue(callable(func))
|
|
self.assertEqual(args[1]['target_dedicated_nodes'], 5)
|
|
self.assertEqual(args[1]['target_low_priority_nodes'], 8)
|
|
return func(*args)
|
|
|
|
mock_pool = mock.create_autospec(models.CloudPool)
|
|
mock_pool.id = "pool id"
|
|
self.mock_self.pools = [mock_pool]
|
|
|
|
self.mock_self._call = call
|
|
self.mock_self.selected_pool = mock.Mock(index=0)
|
|
self.mock_self.ui = mock.create_autospec(PoolsUI)
|
|
|
|
AzureBatchPools.resize_pool(self.mock_self, 5, 8)
|
|
self.assertFalse(self.mock_self.ui.refresh.call_count)
|
|
|
|
AzureBatchPools.resize_pool(self.mock_self, "5", "8")
|
|
self.assertFalse(self.mock_self.ui.refresh.call_count)
|
|
|
|
AzureBatchPools.resize_pool(self.mock_self, None, None)
|
|
self.assertEqual(self.mock_self.ui.refresh.call_count, 1)
|