227 строки
12 KiB
Python
227 строки
12 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
|
|
|
|
import azurebatchutils
|
|
from azurebatchutils import ProgressBar
|
|
from ui_submission import SubmissionUI
|
|
from ui_shared import AzureBatchUI
|
|
from submission import AzureBatchSubmission, AzureBatchRenderJob
|
|
from assets import AzureBatchAssets
|
|
from pools import AzureBatchPools
|
|
from environment import AzureBatchEnvironment
|
|
from shared import AzureBatchSettings
|
|
from exception import CancellationException
|
|
from azurebatchutils import OperatingSystem
|
|
|
|
from azure.batch_extensions import BatchExtensionsClient
|
|
from azure.batch.batch_auth import SharedKeyCredentials
|
|
from azure.batch_extensions import operations
|
|
from azure.batch_extensions import models
|
|
|
|
from azure.storage.blob import BlockBlobService
|
|
|
|
|
|
def print_status(status):
|
|
print(status)
|
|
|
|
|
|
class TestBatchSubmission(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
test_dir = os.path.dirname(__file__)
|
|
top_dir = os.path.dirname(test_dir)
|
|
src_dir = os.path.join(top_dir, 'azure_batch_maya', 'scripts')
|
|
mod_dir = os.path.join(test_dir, 'data', 'modules')
|
|
ui_dir = os.path.join(src_dir, 'ui')
|
|
tools_dir = os.path.join(src_dir, 'tools')
|
|
os.environ["AZUREBATCH_ICONS"] = os.path.join(top_dir, 'azure_batch_maya', 'icons')
|
|
os.environ["AZUREBATCH_TEMPLATES"] = os.path.join(top_dir, 'azure_batch_maya', 'templates')
|
|
os.environ["AZUREBATCH_MODULES"] = mod_dir
|
|
os.environ["AZUREBATCH_SCRIPTS"] = "{0};{1};{2}".format(src_dir, ui_dir, tools_dir)
|
|
os.environ["AZUREBATCH_VERSION"] = "0.1"
|
|
self.mock_self = mock.create_autospec(AzureBatchSubmission)
|
|
self.mock_self.batch = mock.create_autospec(BatchExtensionsClient)
|
|
self.mock_self.batch.job = mock.create_autospec(operations.ExtendedJobOperations)
|
|
self.mock_self.batch.job.jobparameter_from_json.return_value = \
|
|
mock.create_autospec(models.ExtendedJobParameter)
|
|
self.mock_self.ui = mock.create_autospec(SubmissionUI)
|
|
self.mock_self.ui.submit_status = print_status
|
|
self.mock_self.pool_manager = mock.create_autospec(AzureBatchPools)
|
|
self.mock_self.asset_manager = mock.create_autospec(AzureBatchAssets)
|
|
self.mock_self.env_manager = mock.create_autospec(AzureBatchEnvironment)
|
|
self.mock_self._log = logging.getLogger("TestSubmission")
|
|
self.mock_self.renderer = None
|
|
self.mock_self.frame = mock.create_autospec(AzureBatchUI)
|
|
return super(TestBatchSubmission, self).setUp()
|
|
|
|
@mock.patch.object(AzureBatchSubmission, "_collect_modules")
|
|
@mock.patch("submission.callback")
|
|
@mock.patch("submission.SubmissionUI")
|
|
def test_submission_create(self, mock_ui, mock_call, mock_mods):
|
|
submission = AzureBatchSubmission(2, "frame", "call")
|
|
mock_mods.assert_called_with()
|
|
mock_ui.assert_called_with(submission, "frame")
|
|
#mock_call.after_new.assert_called_with(mock.ANY)
|
|
#mock_call.after_open.assert_called_with(mock.ANY)
|
|
|
|
@mock.patch("submission.maya")
|
|
def test_submission_start(self, mock_maya):
|
|
session = mock.Mock(batch="batch_client", storage="storage_client")
|
|
self.mock_self.ui = mock.create_autospec(SubmissionUI)
|
|
self.mock_self.ui.render_module = "module"
|
|
self.mock_self.renderer = mock.Mock()
|
|
AzureBatchSubmission.start(self.mock_self, session, "assets", "pools", "env")
|
|
self.mock_self.renderer.delete.assert_called_with()
|
|
self.mock_self._configure_renderer.assert_called_with()
|
|
self.mock_self.renderer.display.assert_called_with("module")
|
|
self.mock_self.ui.is_logged_in.assert_called_with()
|
|
|
|
def test_submission_collect_modules(self):
|
|
mods = AzureBatchSubmission._collect_modules(self.mock_self)
|
|
self.assertEqual(len(mods), 4)
|
|
|
|
@mock.patch("submission.utils")
|
|
@mock.patch("submission.AzureBatchRenderJob")
|
|
@mock.patch("submission.maya")
|
|
def test_submission_configure_renderer(self, mock_maya, mock_default, mock_utils):
|
|
mock_default.return_value = mock.Mock(render_engine = "default")
|
|
mock_maya.get_attr.return_value = "test_renderer"
|
|
mock_utils.get_current_scene_renderer.return_value = ""
|
|
|
|
renderer = mock.Mock(render_engine = "my_renderer")
|
|
self.mock_self.modules = [renderer, "test", None]
|
|
|
|
AzureBatchSubmission._configure_renderer(self.mock_self)
|
|
self.assertEqual(self.mock_self.renderer.render_engine, "default")
|
|
|
|
renderer = mock.Mock(render_engine = "test_renderer")
|
|
self.mock_self.modules.append(renderer)
|
|
|
|
mock_utils.get_current_scene_renderer.return_value = "test_renderer"
|
|
AzureBatchSubmission._configure_renderer(self.mock_self)
|
|
self.assertEqual(self.mock_self.renderer.render_engine, "test_renderer")
|
|
|
|
def test_submission_refresh_renderer(self):
|
|
self.mock_self.ui = mock.create_autospec(SubmissionUI)
|
|
self.mock_self.renderer = mock.Mock()
|
|
AzureBatchSubmission.refresh_renderer(self.mock_self, "layout")
|
|
self.mock_self.renderer.delete.assert_called_with()
|
|
self.mock_self.renderer.display.assert_called_with("layout")
|
|
|
|
def test_submission_available_pools(self):
|
|
def list_pools(**kwargs):
|
|
self.assertTrue(kwargs.get("lazy"))
|
|
return ["pool1", "pool2"]
|
|
self.mock_self.pool_manager = mock.Mock(list_pools=list_pools)
|
|
pools = AzureBatchSubmission.available_pools(self.mock_self)
|
|
self.assertEqual(pools, ["pool1", "pool2"])
|
|
|
|
@mock.patch("submission.utils")
|
|
@mock.patch("submission.maya")
|
|
def test_submission_submit(self, mock_maya, mock_utils):
|
|
def call(func, *args, **kwargs):
|
|
self.assertTrue(callable(func))
|
|
return func(*args, **kwargs)
|
|
|
|
mock_prog = mock.create_autospec(ProgressBar)
|
|
mock_prog.is_cancelled.return_value = False
|
|
mock_utils.ProgressBar.return_value = mock_prog
|
|
mock_utils.format_scene_path.return_value = "test_file_path"
|
|
mock_utils.build_template_filename.side_effect = azurebatchutils.build_template_filename
|
|
self.mock_self._configure_pool = lambda t: AzureBatchSubmission._configure_pool(self.mock_self, t)
|
|
self.mock_self._submit_threads = lambda: 6
|
|
self.mock_self._check_plugins.return_value = []
|
|
self.mock_self._get_os_flavor.return_value = OperatingSystem.windows
|
|
self.mock_self.pool_manager.create_auto_pool.return_value = {'autoPool': 'auto-pool'}
|
|
self.mock_self.pool_manager.create_pool.return_value = {'poolId': 'new-pool'}
|
|
self.mock_self.env_manager.get_environment_settings.return_value = [{'name':'foo', 'value':'bar'}]
|
|
self.mock_self.renderer = mock.Mock(render_engine='arnold')
|
|
self.mock_self.renderer.get_jobdata.return_value = ("a", "b")
|
|
self.mock_self.renderer.get_params.return_value = {"foo": "bar"}
|
|
self.mock_self.renderer.get_title.return_value = "job name"
|
|
self.mock_self._get_task_container_image.return_value = "containerImage"
|
|
self.mock_self._call = call
|
|
mock_job = mock.create_autospec(models.ExtendedJobParameter)
|
|
self.mock_self.batch.job.jobparameter_from_json.return_value = mock_job
|
|
self.mock_self.asset_manager.upload.return_value = ({"project":"files", "path_map":"maps", "thumb_script":"thumbs", "workspace":"workspace"}, mock_prog)
|
|
self.mock_self.asset_manager.generate_sas_token.return_value = "0123456789ABCDEF"
|
|
self.mock_self.batch.threads = 6
|
|
mock_maya.about.return_value = "2017"
|
|
|
|
self.mock_self.ui.get_pool.return_value = {1: (4, 4)}
|
|
AzureBatchSubmission.submit(self.mock_self)
|
|
self.assertEqual(mock_maya.error.call_count, 0)
|
|
self.mock_self.renderer.disable.assert_called_with(True)
|
|
self.mock_self.pool_manager.create_auto_pool.assert_called_with((4, 4), "job name")
|
|
self.mock_self.batch.job.add.assert_called_with(mock_job)
|
|
self.mock_self.batch.job.jobparameter_from_json.assert_called_with(
|
|
{'commonEnvironmentSettings': [{'name': 'foo', 'value':'bar'}],
|
|
'poolInfo': {'autoPool': 'auto-pool'},
|
|
'displayName': 'job name',
|
|
'id': mock.ANY,
|
|
'applicationTemplateInfo': {
|
|
'parameters': {'taskContainerImageName' : 'containerImage', 'sceneFile': 'test_file_path', 'outputs': mock.ANY, 'assetScript': 'maps', 'foo': 'bar',
|
|
'projectData': 'files', 'thumbScript': 'thumbs', 'storageURL': '0123456789ABCDEF', 'workspace': 'workspace'},
|
|
'filePath': os.path.join(os.environ['AZUREBATCH_TEMPLATES'], 'containers', 'arnold-2017-windows.json')},
|
|
'metadata': [{'name': 'JobType', 'value': 'Maya'}]})
|
|
|
|
|
|
self.mock_self.ui.get_pool.return_value = {2:4}
|
|
AzureBatchSubmission.submit(self.mock_self)
|
|
self.assertEqual(mock_maya.error.call_count, 0)
|
|
self.mock_self.renderer.disable.assert_called_with(True)
|
|
self.mock_self.batch.job.add.assert_called_with(mock_job)
|
|
self.mock_self.batch.job.jobparameter_from_json.assert_called_with(
|
|
{'commonEnvironmentSettings': [{'name': 'foo', 'value':'bar'}],
|
|
'poolInfo': {'poolId': '4'},
|
|
'displayName': 'job name',
|
|
'id': mock.ANY,
|
|
'applicationTemplateInfo': {
|
|
'parameters': {'taskContainerImageName' : 'containerImage','sceneFile': 'test_file_path', 'outputs': mock.ANY, 'assetScript': 'maps', 'foo': 'bar',
|
|
'projectData': 'files', 'thumbScript': 'thumbs', 'storageURL': '0123456789ABCDEF', 'workspace': 'workspace'},
|
|
'filePath': os.path.join(os.environ['AZUREBATCH_TEMPLATES'], 'containers', 'arnold-2017-windows.json')},
|
|
'metadata': [{'name': 'JobType', 'value': 'Maya'}]})
|
|
|
|
|
|
self.mock_self._check_outputs.side_effect = ValueError("No camera")
|
|
AzureBatchSubmission.submit(self.mock_self)
|
|
self.assertEqual(mock_maya.error.call_count, 1)
|
|
self.mock_self._check_outputs.side_effect = None
|
|
|
|
self.mock_self.ui.get_pool.return_value = {3: (4, 4)}
|
|
AzureBatchSubmission.submit(self.mock_self)
|
|
self.assertEqual(mock_maya.error.call_count, 1)
|
|
self.mock_self.renderer.disable.assert_called_with(True)
|
|
|
|
self.mock_self.batch.job.add.assert_called_with(mock_job)
|
|
self.mock_self.batch.job.add.call_count = 0
|
|
self.mock_self.pool_manager.create_pool.assert_called_with((4, 4), 'job name')
|
|
|
|
mock_prog.is_cancelled.side_effect = CancellationException("cancelled")
|
|
AzureBatchSubmission.submit(self.mock_self)
|
|
self.assertEqual(mock_maya.info.call_count, 4)
|
|
self.mock_self.renderer.disable.assert_called_with(True)
|
|
self.assertEqual(self.mock_self.batch.job.add.call_count, 0)
|
|
|
|
mock_prog.is_cancelled.side_effect = None
|
|
self.mock_self.pool_manager.create_pool.side_effect = ValueError("Bad data")
|
|
AzureBatchSubmission.submit(self.mock_self)
|
|
self.assertEqual(mock_maya.error.call_count, 2)
|
|
self.mock_self.renderer.disable.assert_called_with(True)
|
|
self.assertEqual(self.mock_self.batch.job.add.call_count, 0)
|