"""Unit tests for PISCES tasks."""
import os
import sys

import pytest

from mined_metric.builder.metrics_impl.pisces.task.meta import (
    _create_sdl_pipedream_job_json,
    _run_sdl_pipedream_pretend,
    SteeringMode,
)
from mined_metric.builder.metrics_impl.pisces.task.standard import (
    PiscesStandardTask,
)
from vehicle.planner.metrics.pipegraph import PipeGraph


def test_run_sdl_pipedream_pretend():
    """Tests creating a simulation PipeGraph."""
    json = _create_sdl_pipedream_job_json("id", "filename")
    pipegraph = _run_sdl_pipedream_pretend(
        PipeGraph("name"),
        "user@zoox.com",
        "sha",
        "label",
        json,
        SteeringMode.k2WS,
    )
    job = pipegraph.render()[0]
    assert "set -euxo pipefail" in job.command
    assert (
        "mkdir -p /mnt/nautilus_rw/pisces/sdl_pipedream_json/id" in job.command
    )
    assert "sim/launch.sh pipedream planner" in job.command
    assert (
        "//mined_metric/builder/metrics_impl/pisces:steering_mode_log_tests"
        in job.command
    )
    assert '--simulator_args="' in job.command
    assert (
        "--params-kv sim/override_autonomous_mode=PURE_PLAYBACK_MODE"
        in job.command
    )
    assert "--params-kv sim/vehicle_type=vh6" in job.command
    assert "--params-kv sim/sim_model=vh6a" in job.command
    assert "--params-kv sim/terminate_on=none" in job.command
    assert '--params-kv planner/config_variant_id=vh6_2ws"' in job.command
    assert "--alert_email user@zoox.com" in job.command
    assert "--partition=main_turing" in job.command
    assert "--labels='[\"label\"]'" in job.command
    assert "--pretend" in job.command


@pytest.mark.parametrize("tracking", [None, "test_tracking"])
def test_tracking_arg(tracking):
    """Test specifying a tracking branch."""
    pipegraph = PipeGraph("name")
    PiscesStandardTask.add_tasks(
        pipegraph,
        validation_id="",
        num_workers=0,
        user_email="",
        control_sha="0123456789abcdef0123456789abcdef01234567",
        control_label="",
        candidate_sha="0123456789abcdef0123456789abcdef01234567",
        candidate_label="",
        tracking=tracking,
        notification_target=None,
        run_meta_ids=None,
        dataset_path=None,
        analysis_branch=None,
        with_pcp=False,
        sdl_pipedream_control_2ws_json_path=None,
        sdl_pipedream_control_4ws_json_path=None,
        sdl_pipedream_candidate_2ws_json_path=None,
        sdl_pipedream_candidate_4ws_json_path=None,
        steering_mode_comparison_only=False,
    )
    # Check whether --tracking argument was added to "analyze" job.
    analyze = pipegraph.pipedefs["analyze"]
    if tracking:
        assert f"--tracking {tracking}" in analyze._data["arg"]
    else:
        assert all(["--tracking" not in arg for arg in analyze._data["arg"]])


if __name__ == "__main__":
    args = [os.path.dirname(__file__), "--color=yes", "--verbose", "--verbose"]
    sys.exit(pytest.main(args))