taskcluster/taskgraph/test/test_decision.py
author Dustin J. Mitchell <dustin@mozilla.com>
Wed, 18 May 2016 18:02:51 +0000
changeset 298016 e5349d525d9d3e60e6992d619a6cf56a270051e8
parent 297716 c3e24e94ab5148ad92a3f137d6cab9d8bebdd6a7
child 300760 051b7c10b6175d6aba58af013f8e4e289c6d05d7
permissions -rw-r--r--
Bug 1273673: use Python logging for taskgraph generation; r=gps MozReview-Commit-ID: EYI6hGFUWP8

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from __future__ import absolute_import, print_function, unicode_literals

import os
import json
import yaml
import shutil
import unittest
import tempfile

from .. import decision
from ..graph import Graph
from ..types import Task, TaskGraph
from mozunit import main

class TestDecision(unittest.TestCase):

    def test_taskgraph_to_json(self):
        tasks = {
            'a': Task(kind=None, label='a', attributes={'attr': 'a-task'}),
            'b': Task(kind=None, label='b', task={'task': 'def'}),
        }
        graph = Graph(nodes=set('ab'), edges={('a', 'b', 'edgelabel')})
        taskgraph = TaskGraph(tasks, graph)

        res = decision.taskgraph_to_json(taskgraph)

        self.assertEqual(res, {
            'a': {
                'attributes': {'attr': 'a-task'},
                'task': {},
                'dependencies': [('edgelabel', 'b')],
            },
            'b': {
                'attributes': {},
                'task': {'task': 'def'},
                'dependencies': [],
            }
        })


    def test_write_artifact_json(self):
        data = [{'some': 'data'}]
        tmpdir = tempfile.mkdtemp()
        try:
            decision.ARTIFACTS_DIR = os.path.join(tmpdir, "artifacts")
            decision.write_artifact("artifact.json", data)
            with open(os.path.join(decision.ARTIFACTS_DIR, "artifact.json")) as f:
                self.assertEqual(json.load(f), data)
        finally:
            if os.path.exists(tmpdir):
                shutil.rmtree(tmpdir)
            decision.ARTIFACTS_DIR = 'artifacts'


    def test_write_artifact_yml(self):
        data = [{'some': 'data'}]
        tmpdir = tempfile.mkdtemp()
        try:
            decision.ARTIFACTS_DIR = os.path.join(tmpdir, "artifacts")
            decision.write_artifact("artifact.yml", data)
            with open(os.path.join(decision.ARTIFACTS_DIR, "artifact.yml")) as f:
                self.assertEqual(yaml.safe_load(f), data)
        finally:
            if os.path.exists(tmpdir):
                shutil.rmtree(tmpdir)
            decision.ARTIFACTS_DIR = 'artifacts'


if __name__ == '__main__':
    main()