
from collections import defaultdict
import json
import os

from pysiv2.custom.base import TestValuesLoader


class TestChunkedTasks(TestValuesLoader):
    """
    Sanity checks related to pbsmrtpipe chunking (worfklow-level
    parallelization).
    """

    def setUp(self):
        self._was_chunked = False
        workflow_options = os.path.join(self.job_dir, "workflow",
                                        "options-workflow.json")
        with open(workflow_options) as f:
            options = json.load(f)
            self._was_chunked = options["pbsmrtpipe.options.chunk_mode"]

    def test_task_chunking(self):
        """
        Simple test to verify that any tasks we expect to be chunked are
        actually chunked.  These should be explicitly listed by task ID in
        ``test_values.json`` using the key ``chunked_tasks``, e.g.::

            "chunked_tasks": ["pbccs.tasks.ccs"]
        """
        task_ids = self.test_values["chunked_tasks"]
        tasks_dir = os.path.join(self.job_dir, "tasks")
        tasks_dict = defaultdict(list)
        for task_dir in os.listdir(tasks_dir):
            if task_dir.startswith("."):
                continue
            task_id, job_number = task_dir.split("-")
            tasks_dict[task_id].append(int(job_number))
        for task_id in task_ids:
            if self._was_chunked:
                self.assertTrue(tasks_dict[task_id] != [0],
                    "Task {t} was not chunked".format(t=task_id))
            else:
                self.assertTrue(tasks_dict[task_id] == [0])
