# -*- encoding: utf-8 -*- import responses import pillarsdk import pillar.tests import pillar.auth import pillar.tests.common_test_data as ctd from abstract_attract_test import AbstractAttractTest class ShotManagerTest(AbstractAttractTest): def setUp(self, **kwargs): AbstractAttractTest.setUp(self, **kwargs) self.tmngr = self.app.pillar_extensions['attract'].task_manager self.smngr = self.app.pillar_extensions['attract'].shot_manager self.proj_id, self.project = self.ensure_project_exists() self.sdk_project = pillarsdk.Project(pillar.tests.mongo_to_sdk(self.project)) def create_task(self, shot_id, task_type): with self.app.test_request_context(): # Log in as project admin user pillar.auth.login_user(ctd.EXAMPLE_PROJECT_OWNER_ID) self.mock_blenderid_validate_happy() task = self.tmngr.create_task(self.sdk_project, parent=shot_id, task_type=task_type) self.assertIsInstance(task, pillarsdk.Node) return task def create_shot(self): with self.app.test_request_context(): # Log in as project admin user pillar.auth.login_user(ctd.EXAMPLE_PROJECT_OWNER_ID) self.mock_blenderid_validate_happy() shot = self.smngr.create_shot(self.sdk_project) self.assertIsInstance(shot, pillarsdk.Node) return shot @responses.activate def test_tasks_for_shot(self): shot1 = self.create_shot() shot2 = self.create_shot() shot1_id = shot1['_id'] shot2_id = shot2['_id'] task1 = self.create_task(shot1_id, u'fx') task2 = self.create_task(shot1_id, u'fx') task3 = self.create_task(shot1_id, u'høken') task4 = self.create_task(shot2_id, u'effects') task5 = self.create_task(shot2_id, u'effects') task6 = self.create_task(shot2_id, u'ïnžane') with self.app.test_request_context(): # Log in as project admin user pillar.auth.login_user(ctd.EXAMPLE_PROJECT_OWNER_ID) self.mock_blenderid_validate_happy() shot_id_to_task = self.smngr.tasks_for_shots([shot1, shot2], [u'fx', u'høken', u'effects']) # Just test based on task IDs, as strings are turned into datetimes etc. by the API, # so we can't test equality. for all_tasks in shot_id_to_task.values(): for task_type, tasks in all_tasks.items(): all_tasks[task_type] = {task['_id'] for task in tasks} self.assertEqual({ u'fx': {task1['_id'], task2['_id']}, u'høken': {task3['_id']}, }, shot_id_to_task[shot1_id]) self.assertEqual({ u'effects': {task4['_id'], task5['_id']}, None: {task6['_id']}, }, shot_id_to_task[shot2_id])