test_snakemake.py

from collections import defaultdict
from pathlib import Path
from typing import Any
from typing import Callable
from typing import Dict
from typing import List
from typing import Optional

import snakemake

class SnakemakeLogger:
    """Returns a log handler for snakemake to track how many times a rule was run.

    Attributes:
       counter: the mapping from rule name to the # of times that rule was run.
       log_handler: the log handler to provide to snakemake
    """

    def __init__(self) -> None:
        self.counter: Dict[str, int] = defaultdict(lambda: 0)
        self.log_handler: Callable[[Dict[str, Any]], None]

        def log_handler(d: Dict[str, Any]) -> None:
            if d["level"] != "run_info":
                return
            # NB: skip the first two and last lines
            for counts_line in d["msg"].split("\n")[2:-1]:
                counts_line = counts_line.strip()
                count, job = counts_line.split("\t")
                assert int(count) > 0, counts_line
                self.counter[job] += int(count)

        self.log_handler = log_handler


def run_snakemake(
    snakefile: Path,
    workdir: Path,
    resources: Optional[Dict[str, Any]] = None,
    config: Optional[Dict[str, Any]] = None,
    configfiles: Optional[List[Path]] = None,
) -> SnakemakeLogger:
    """Runs Snakemake in dry-run mode.
    
    Tests that snakefile can be parsed by Snakemake.
    
    Returns the `SnakemakeLogger` so that the caller may verify which rules are run
    how many times.
    
    Any input data required by the workflow needs to be mocked prior to calling this
    function.

    Args:
        snakefile: the snake file to execute
        workdir: the working directory in which to run Snakemake
        resources: the mapping of resource identifier to value (eg. `{"mem_gb": 8}`)
        config: the configuration object for Snakemake
        configfiles: one or more config files to provide to Snakemake
    """
    assert snakefile.is_file(), f"{snakefile} is not a file"

    logger: SnakemakeLogger = SnakemakeLogger()
    
    assert snakemake.snakemake(
        snakefile=str(snakefile),
        config=config,
        configfiles=configfiles,
        resources=resources,
        workdir=str(workdir),
        dryrun=True,
        quiet=True,
        log_handler=[logger.log_handler],
        ignore_ambiguity=True,
    ), f"Snakemake failed: {snakefile}"

    assert logger.rule_count["all"] == 1

    return logger