import datetime
import json
import os
import responses
import pyarrow as pa
import base64
from unittest import mock

import hashlib
import wallaroo
from wallaroo import expression as e
from wallaroo.pipeline_config import *
from wallaroo.model_version import ModelVersion
from . import testutil
from .reusable_responders import add_insert_model_config_response_with_config


# TODO: derive this test class from unitttest.TestCase so we can use self.assertXXX etc


class TestPipelineConfig:
    def setup_method(self):
        self.ix = 0
        self.now = datetime.datetime.now()
        self.gql_client = testutil.new_gql_client(
            endpoint="http://api-lb:8080/v1/graphql"
        )
        self.test_client = wallaroo.Client(
            gql_client=self.gql_client, auth_type="test_auth"
        )

    def test_simple(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        assert pipeline.config() == PipelineConfig("test", [], [])

    def gen_id(self):
        self.ix += 1
        return self.ix

    def ccfraud_model(self, variant="baseline", config={}):
        model_config = {"runtime": "onnx", **config}

        fake_data = {
            "model_id": f"ccfraud_{variant}",
            "model_version": variant,
            "sha": "ccfraud_sha",
            "file_name": f"ccfraud_{variant}.onnx",
        }

        return self.fake_model(model_config, fake_data)

    @responses.activate
    def fake_model(self, model_config, fake_data):
        add_insert_model_config_response_with_config(
            self.test_client.api_endpoint, self.gen_id(), model_config
        )

        model = ModelVersion(
            self.test_client,
            data={
                "id": self.gen_id(),
                "sha": hashlib.sha256(
                    (json.dumps(model_config) + json.dumps(fake_data)).encode("utf8")
                ).hexdigest(),
                "updated_at": self.now.isoformat(),
                **fake_data,
            },
        )

        model.configure(**model_config)

        # sorry
        model._config._model_version = model

        return model

    @responses.activate
    @mock.patch.dict(os.environ, {"MODELS_ENABLED": "true"})
    def mlflow_model(self, input_schema, output_schema, variant="primary"):
        upload_resp = responses.Response(
            method=responses.POST,
            url="http://api-lb:8080/v1/api/models/upload",
            status=200,
            json={"insert_models": {"returning": [{"models": [{"id": 1}]}]}},
        )
        responses.add(upload_resp)

        udw_resp = responses.Response(
            method=responses.POST,
            url="http://api-lb:8080/v1/graphql",
            status=200,
            match=[testutil.query_name_matcher("UserDefaultWorkspace")],
            json={
                "data": {
                    "user_default_workspace": [
                        {
                            "workspace": {
                                "id": 1,
                            }
                        }
                    ]
                }
            },
        )
        responses.add(udw_resp)

        input_schema_str = base64.b64encode(bytes(input_schema.serialize())).decode(
            "utf8"
        )
        output_schema_str = base64.b64encode(bytes(output_schema.serialize())).decode(
            "utf8"
        )

        get_version_resp = responses.Response(
            method=responses.POST,
            url=f"{self.test_client.api_endpoint}/v1/api/models/get_version_by_id",
            status=200,
            json={
                "model_version": {
                    "model_version": {
                        "name": "new-model",
                        "visibility": "private",
                        "workspace_id": 1,
                        "conversion": {
                            "python_version": "3.8",
                            "requirements": [],
                            "framework": "keras",
                        },
                        "id": 1,
                        "image_path": None,
                        "status": "ready",
                        "task_id": "7f05c403-dcf4-4ecb-b5ea-28f27aa7eb7b",
                        "file_info": {
                            "version": "ec1ab8e3-923b-40dd-9f77-f20bbe8058b3",
                            "sha": "f7e49538e38bebe066ce8df97bac8be239ae8c7d2733e500c8cd633706ae95a8",
                            "file_name": "simple_model.h5",
                        },
                    },
                    "config": {
                        "id": 1,
                        "model_version_id": 1,
                        "runtime": "mlflow",
                        "input_schema": input_schema_str,
                        "output_schema": output_schema_str,
                    },
                }
            },
        )
        responses.add(get_version_resp)

        insert_config_resp = responses.Response(
            responses.POST,
            f"{self.test_client.api_endpoint}/v1/api/models/insert_model_config",
            status=200,
            json={
                "model_config": {
                    "id": 1,
                    "model_version_id": 1,
                    "runtime": "mlflow",
                    "input_schema": input_schema_str,
                    "output_schema": output_schema_str,
                },
            },
        )
        responses.add(insert_config_resp)

        responses.add(
            responses.POST,
            "http://api-lb:8080/v1/graphql",
            status=200,
            match=[testutil.json_body_printer()],
        )

        model = self.test_client.register_model_image(
            "mlflow-primary", "my-image"
        ).configure(
            "mlflow",
            input_schema=input_schema,
            output_schema=output_schema,
        )

        model._fill(
            {
                "id": 3,
                "model_id": "mlflow-primary",
                "model_version": "n/a",
                "file_name": "none",
                "image_path": "my-image",
                "sha": "sha_sha",
                "updated_at": self.now.isoformat(),
            }
        )
        model._config._fill(
            {
                "id": model._config.id,
                "input_schema": model._config._input_schema,
                "output_schema": model._config._output_schema,
                "runtime": model._config._runtime,
                "model": model._config._model_version,
                "batch_config": None,
                "tensor_fields": None,
                "filter_threshold": None,
            }
        )
        return model

    def assertJsonRoundTrip(self, pipeline: PipelineConfigBuilder):
        config = pipeline.config()
        assert PipelineConfig.from_json(config.to_json()) == config

    def assertModelConfigs(
        self, pipeline: PipelineConfigBuilder, models: Iterable[ModelVersion]
    ):
        active = set(mc.id() for mc in pipeline._model_configs())
        for model in models:
            assert model.config().id() in active

    @responses.activate
    def test_add_models(self):
        responses.add(
            responses.POST,
            "http://api-lb:8080/v1/graphql",
            status=200,
            match=[testutil.mutation_name_matcher("ConfigureModel")],
            json={
                "data": {
                    "insert_model_config": {
                        "returning": [
                            {
                                "id": 2,
                            }
                        ]
                    }
                }
            },
        )

        responses.add(
            responses.POST,
            "http://api-lb:8080/v1/graphql",
            status=200,
            match=[testutil.query_name_matcher("ModelConfigById")],
            json={
                "data": {
                    "model_config_by_pk": {
                        "id": 1,
                        "filter_threshold": 0.1234,
                        "model": {
                            "id": 2,
                        },
                        "runtime": "onnx",
                        "tensor_fields": None,
                    },
                },
            },
        )

        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")

        pipeline.add_multi_model_step([one, two])
        assert pipeline.steps == [
            ModelInference(
                [
                    ModelForStep.from_model(one),
                    ModelForStep.from_model(two),
                ]
            )
        ]
        self.assertModelConfigs(pipeline, [one, two])
        self.assertJsonRoundTrip(pipeline)

    def test_remove_step(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")

        pipeline.add_model_step(one)
        pipeline.add_multi_model_step([one, two])
        pipeline.remove_step(0)

        assert len(pipeline.model_configs) == 1
        assert pipeline.steps == [
            ModelInference(
                [
                    ModelForStep.from_model(one),
                    ModelForStep.from_model(two),
                ]
            )
        ]

    def test_replace_with_model_step(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")

        pipeline.add_model_step(one)
        pipeline.add_multi_model_step([one, two])
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference(
                [
                    ModelForStep.from_model(one),
                    ModelForStep.from_model(two),
                ]
            ),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 3
        pipeline.replace_with_model_step(1, two)

        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference([ModelForStep.from_model(two)]),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 2

    def test_replace_with_multi_model_step(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")

        pipeline.add_model_step(one)
        pipeline.add_model_step(two)
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference([ModelForStep.from_model(two)]),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 2
        pipeline.replace_with_multi_model_step(0, [one, two])

        assert pipeline.steps == [
            ModelInference(
                [
                    ModelForStep.from_model(one),
                    ModelForStep.from_model(two),
                ]
            ),
            ModelInference([ModelForStep.from_model(two)]),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 3

    def test_audit(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")

        pipeline.add_multi_model_step([one]).add_audit("1")
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            AuditResults(1, None),
        ]
        self.assertModelConfigs(pipeline, [one])
        self.assertJsonRoundTrip(pipeline)

    def test_replace_with_audit(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")
        pipeline.add_model_step(one)
        pipeline.add_model_step(two)
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference([ModelForStep.from_model(two)]),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 2

        pipeline.replace_with_audit(1, "1")
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 1
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            AuditResults(1, None),
        ]

    def test_key_split(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("baseline")
        gold_model = self.ccfraud_model("gold")
        pipeline.add_key_split(baseline, "card_type", {"gold": gold_model})
        assert pipeline.steps == [
            MetaValueSplit(
                "card_type",
                ModelForStep.from_model(baseline),
                {"gold": ModelForStep.from_model(gold_model)},
            )
        ]
        self.assertModelConfigs(pipeline, [baseline, gold_model])
        self.assertJsonRoundTrip(pipeline)

    def test_replace_with_key_split(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")
        pipeline.add_model_step(one)
        pipeline.add_model_step(two)
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference([ModelForStep.from_model(two)]),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 2

        baseline = self.ccfraud_model("baseline")
        gold_model = self.ccfraud_model("gold")
        pipeline.replace_with_key_split(1, baseline, "card_type", {"gold": gold_model})
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            MetaValueSplit(
                "card_type",
                ModelForStep.from_model(baseline),
                {"gold": ModelForStep.from_model(gold_model)},
            ),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 3

    def test_random_split(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("a")
        bandit = self.ccfraud_model("b")
        pipeline.add_random_split([(95, baseline), (5, bandit)])
        assert pipeline.steps == [
            RandomSplit(
                [
                    ModelWeight(95, ModelForStep.from_model(baseline)),
                    ModelWeight(5, ModelForStep.from_model(bandit)),
                ]
            )
        ]
        self.assertModelConfigs(pipeline, [baseline, bandit])
        self.assertJsonRoundTrip(pipeline)

    def test_replace_with_random_split(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")
        pipeline.add_model_step(one)
        pipeline.add_model_step(two)
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference([ModelForStep.from_model(two)]),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 2

        baseline = self.ccfraud_model("a")
        bandit = self.ccfraud_model("b")
        pipeline.replace_with_random_split(0, [(95, baseline), (5, bandit)])
        assert pipeline.steps == [
            RandomSplit(
                [
                    ModelWeight(95, ModelForStep.from_model(baseline)),
                    ModelWeight(5, ModelForStep.from_model(bandit)),
                ]
            ),
            ModelInference([ModelForStep.from_model(two)]),
        ]
        assert len(pipeline.model_configs) == 2
        assert len(pipeline._model_configs()) == 3

    def test_shadow_deploy(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("baseline")
        burn_in = self.ccfraud_model("burn_in")
        pipeline.add_shadow_deploy(baseline, [burn_in])
        assert pipeline.steps == [
            ModelInference(
                [
                    ModelForStep.from_model(baseline),
                    ModelForStep.from_model(burn_in),
                ]
            ),
            AuditResults(1),
            MultiOut(),
        ]
        self.assertModelConfigs(pipeline, [baseline, burn_in])
        self.assertJsonRoundTrip(pipeline)

    def test_replace_with_shadow_deploy(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")
        three = self.ccfraud_model("three")
        pipeline.add_model_step(one)
        pipeline.add_model_step(two)
        pipeline.add_model_step(three)
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference([ModelForStep.from_model(two)]),
            ModelInference([ModelForStep.from_model(three)]),
        ]
        assert len(pipeline.model_configs) == 3
        assert len(pipeline._model_configs()) == 3

        baseline = self.ccfraud_model("baseline")
        burn_in = self.ccfraud_model("burn_in")
        pipeline.replace_with_shadow_deploy(1, baseline, [burn_in])
        assert pipeline.steps == [
            ModelInference([ModelForStep.from_model(one)]),
            ModelInference(
                [
                    ModelForStep.from_model(baseline),
                    ModelForStep.from_model(burn_in),
                ]
            ),
            AuditResults(1),
            MultiOut(),
            ModelInference([ModelForStep.from_model(three)]),
        ]
        assert len(pipeline.model_configs) == 5
        assert len(pipeline._model_configs()) == 4

    def test_simple_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        ccfraud = self.ccfraud_model("keras")

        pipeline.add_model_step(ccfraud)

        self.assertModelConfigs(pipeline, [ccfraud])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(ccfraud._config.to_k8s_yaml(), "model_config.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_batch_config_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        ccfraud = self.ccfraud_model("keras", config={"batch_config": "single"})

        pipeline.add_model_step(ccfraud)

        self.assertModelConfigs(pipeline, [ccfraud])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(ccfraud._config.to_k8s_yaml(), "model_config.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_py_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")

        wl_model = self.fake_model(
            {"runtime": "onnx"},
            {
                "model_id": f"demandcurve",
                "model_version": "0",
                "file_name": f"demand_curve_made_here.onnx",
            },
        )

        model_config = {
            "runtime": "python",
        }
        model_pre = self.fake_model(
            model_config,
            {
                "model_id": "preprocess",
                "model_version": "0",
                "file_name": "day5_preprocess.py",
            },
        )
        model_post = self.fake_model(
            model_config,
            {
                "model_id": "postprocess",
                "model_version": "0",
                "file_name": "day5_postprocess.py",
            },
        )

        pipeline = (
            pipeline.add_model_step(model_pre)
            .add_model_step(wl_model)
            .add_model_step(model_post)
        )

        self.assertModelConfigs(pipeline, [model_pre, wl_model, model_post])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(model_pre._config.to_k8s_yaml(), "pre_model_config.yaml")
        snapshot.assert_match(wl_model._config.to_k8s_yaml(), "wl_model_config.yaml")
        snapshot.assert_match(
            model_post._config.to_k8s_yaml(), "post_model_config.yaml"
        )
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_ab_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")

        keras_file = "keras_ccfraud.onnx"
        default_model = self.fake_model(
            {"runtime": "onnx"},
            {
                "model_id": f"keras_ccfraud",
                "model_version": "0",
                "file_name": keras_file,
            },
        )

        experiment_file = "modelA.onnx"
        experiment_model = self.fake_model(
            {"runtime": "onnx"},
            {
                "model_id": "experiment",
                "model_version": "0",
                "file_name": experiment_file,
            },
        )

        meta_key_name = "card_type"
        pipeline = pipeline.add_key_split(
            default_model, meta_key_name, {"gold": experiment_model}
        )

        self.assertModelConfigs(pipeline, [default_model, experiment_model])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(
            default_model._config.to_k8s_yaml(), "default_model_config.yaml"
        )
        snapshot.assert_match(
            experiment_model._config.to_k8s_yaml(), "experiment_model_config.yaml"
        )
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_random_split_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")

        # due to how the engine handles tensorflow and local files, the model
        # name must include .zip and be identical to the file name :sadlol:
        control = self.fake_model(
            {"runtime": "tensorflow"},
            {
                "model_id": f"aloha-cnn-lstm.zip",
                "model_version": "0",
                "file_name": "aloha-cnn-lstm.zip",
            },
        )
        challenger = self.fake_model(
            {"runtime": "tensorflow"},
            {
                "model_id": f"aloha-cnn-lstm-new.zip",
                "model_version": "0",
                "file_name": "aloha-cnn-lstm-new.zip",
            },
        )
        pipeline = pipeline.add_random_split(
            [(2, control), (1, challenger)], "session_id"
        )

        self.assertModelConfigs(pipeline, [control, challenger])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(control._config.to_k8s_yaml(), "control.yaml")
        snapshot.assert_match(challenger._config.to_k8s_yaml(), "challenger.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_shadow_deploy_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")

        keras_file = "keras_ccfraud.onnx"
        champion = self.fake_model(
            {"runtime": "onnx"},
            {
                "model_id": f"keras_ccfraud",
                "model_version": "0",
                "file_name": keras_file,
            },
        )

        keras_file = "keras_xgboost.onnx"
        model2 = self.fake_model(
            {"runtime": "onnx"},
            {
                "model_id": f"ccfraudxgb",
                "model_version": "0",
                "file_name": keras_file,
            },
        )

        experiment_file = "modelA.onnx"
        model3 = self.fake_model(
            {"runtime": "onnx"},
            {
                "model_id": "ccfraudrf",
                "model_version": "0",
                "file_name": experiment_file,
            },
        )

        pipeline.add_shadow_deploy(champion, [model2, model3])

        self.assertModelConfigs(pipeline, [champion, model2, model3])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(champion._config.to_k8s_yaml(), "control.yaml")
        snapshot.assert_match(model2._config.to_k8s_yaml(), "model2.yaml")
        snapshot.assert_match(model3._config.to_k8s_yaml(), "model3.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_mlflow_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        input_schema = pa.schema(
            [
                pa.field("temp", pa.float32()),
                pa.field("holiday", pa.uint8()),
                pa.field("workingday", pa.uint8()),
                pa.field("windspeed", pa.float32()),
            ]
        )

        output_schema = pa.schema([pa.field("sum", pa.float32())])
        mlflow = self.mlflow_model(input_schema, output_schema)

        pipeline.add_model_step(mlflow)

        self.assertModelConfigs(pipeline, [mlflow])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(mlflow.config().to_k8s_yaml(), "model_config.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_mlflow_variable_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        input_schema = pa.schema(
            [
                pa.field("vec0", pa.list_(pa.float32())),
                pa.field("vec1", pa.list_(pa.float32())),
            ]
        )

        output_schema = pa.schema([pa.field("matrix", pa.list_(pa.float32()))])
        mlflow = self.mlflow_model(input_schema, output_schema)

        pipeline.add_model_step(mlflow)

        self.assertModelConfigs(pipeline, [mlflow])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(mlflow._config.to_k8s_yaml(), "model_config.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_mlflow_multidim_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        input_schema = pa.schema(
            [
                pa.field("img", pa.list_(pa.list_(pa.list_(pa.float32())))),
            ]
        )

        output_schema = pa.schema(
            [pa.field("bounds", pa.list_(pa.list_(pa.uint32(), 4)))]
        )
        mlflow = self.mlflow_model(input_schema, output_schema)

        pipeline.add_model_step(mlflow)

        self.assertModelConfigs(pipeline, [mlflow])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(mlflow._config.to_k8s_yaml(), "model_config.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    def test_demandgen_mlflow_pipeline(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        input_schema = pa.schema(
            [
                pa.field("site_id", pa.string()),
                pa.field("channel_id", pa.string()),
                pa.field("product_id", pa.string()),
                pa.field("v0", pa.float64()),
                pa.field("v1", pa.float64()),
                pa.field("v2", pa.float64()),
                pa.field("v3", pa.float64()),
                pa.field("v4", pa.float64()),
                pa.field("v5", pa.float64()),
                pa.field("v6", pa.float64()),
                pa.field("v7", pa.float64()),
                pa.field("v8", pa.float64()),
                pa.field("v9", pa.float64()),
                pa.field("v10", pa.float64()),
                pa.field("v11", pa.float64()),
                pa.field("v12", pa.float64()),
                pa.field("v13", pa.float64()),
                pa.field("v14", pa.float64()),
                pa.field("v15", pa.float64()),
                pa.field("v16", pa.float64()),
                pa.field("v17", pa.float64()),
                pa.field("v18", pa.float64()),
                pa.field("v19", pa.float64()),
            ]
        )

        output_schema = pa.schema(
            [
                pa.field("site_id", pa.string()),
                pa.field("channel_id", pa.string()),
                pa.field("product_id", pa.string()),
                pa.field("v0", pa.list_(pa.float64(), 6)),
                pa.field("outcome", pa.list_(pa.float64(), 6)),
            ]
        )
        mlflow = self.mlflow_model(input_schema, output_schema)

        pipeline.add_model_step(mlflow)

        self.assertModelConfigs(pipeline, [mlflow])
        self.assertJsonRoundTrip(pipeline)

        snapshot.assert_match(mlflow._config.to_k8s_yaml(), "model_config.yaml")
        snapshot.assert_match(pipeline.config().to_yaml(), "pipeline_config.yaml")

    @responses.activate
    def test_validate(self, snapshot):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("a")
        bandit = self.ccfraud_model("b")
        baseline_ok = baseline.config().outputs[0] > 0.95
        bandit_ok = bandit.config().outputs[0] > 0.95
        (
            pipeline.add_random_split([(95, baseline), (5, bandit)])
            .add_validation("bandit_ok", bandit_ok)
            .add_validation("baseline_ok", baseline_ok)
        )

        self.assertModelConfigs(pipeline, [baseline, bandit])
        self.assertJsonRoundTrip(pipeline)
        snapshot.assert_match(
            json.dumps(pipeline.config().to_json(), indent=4), "single_validation.json"
        )

    def test_validation_placement(self, snapshot):
        p = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("base")
        shadow = self.ccfraud_model("shadow")
        better = shadow.config().outputs[0][0] > baseline.config().outputs[0][0]
        p = p.add_shadow_deploy(baseline, [shadow])
        p = p.add_validation("shadow_better", better)

        self.assertModelConfigs(p, [baseline, shadow])
        self.assertJsonRoundTrip(p)
        snapshot.assert_match(
            json.dumps(p.config().to_json(), indent=4), "checked_pipeline.json"
        )

    def test_single_validation(self, snapshot):
        p = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("base")
        bounds = baseline.config().outputs[0][0] > 0.0
        p = p.add_model_step(baseline)
        p = p.add_validation("in_bounds", bounds)

        self.assertModelConfigs(p, [baseline])
        self.assertJsonRoundTrip(p)
        snapshot.assert_match(
            json.dumps(p.config().to_json(), indent=4), "single_validation.json"
        )

    def test_single_alert(self, snapshot):
        p = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("base")
        high_fraud = baseline.config().outputs[0][0] > 0.95
        p = p.add_model_step(baseline)
        p = p.add_alert("high_fraud", e.count(high_fraud, "1h") > 100, [])

        self.assertModelConfigs(p, [baseline])
        self.assertJsonRoundTrip(p)
        snapshot.assert_match(json.dumps(p.config().to_json(), indent=4), "alert.json")
        assert (
            p.alert_configurations[0].expression
            == "sum by (pipeline_id) (pipeline_gauge:high_fraud:left) > 100"
        )

    def test_replace_with_alert_failure(self):
        p = PipelineConfigBuilder(None, "test", "test")
        p.add_select(1)
        assert len(p.model_configs) == 1
        assert len(p._model_configs()) == 0
        baseline = self.ccfraud_model("base")
        p.add_model_step(baseline)
        assert len(p.model_configs) == 2
        assert len(p._model_configs()) == 1
        high_fraud = baseline.config().outputs[0][0] > 0.95
        try:
            p.replace_with_alert(1, "high_fraud", e.count(high_fraud, "1h") > 100, [])
            assert False
        except RuntimeError as er:
            assert True

    def test_replace_with_alert(self):
        p = PipelineConfigBuilder(None, "test", "test")
        baseline = self.ccfraud_model("base")
        p.add_model_step(baseline)
        p.add_model_step(baseline)
        assert len(p.model_configs) == 2
        assert len(p._model_configs()) == 2
        high_fraud = baseline.config().outputs[0][0] > 0.95
        p.replace_with_alert(1, "high_fraud", e.count(high_fraud, "1h") > 100, [])
        assert len(p.model_configs) == 2
        assert len(p._model_configs()) == 1
        assert (
            str(p.steps)
            == """[{'ModelInference': {'models': [{'name': 'ccfraud_base', 'version': 'base', 'sha': 'ccfraud_sha'}]}}, {'Check': {'tree': ['{"values": {"high_fraud:left": {"root": {"node": "fn", "fn": "count", "arguments": [{"node": "binop", "op": ">", "left": {"node": "variable", "variant_id": {"name": "ccfraud_base"}, "position": "output", "key": [0, 0]}, "right": {"node": "literal", "float": 0.95}}, {"node": "literal", "timedelta": "1h"}]}, "required_data": [{"name": "ccfraud_base"}]}}, "gauges": ["high_fraud:left"], "validations": []}']}}]"""
        )

    @responses.activate
    def test_clear(self):
        pipeline = PipelineConfigBuilder(None, "test", "test")
        one = self.ccfraud_model("one")
        two = self.ccfraud_model("two")
        pipeline.add_multi_model_step([one, two])

        result = pipeline.clear()
        self.assertJsonRoundTrip(pipeline)
        assert isinstance(result, PipelineConfigBuilder)
        assert pipeline.steps == []
        self.assertModelConfigs(pipeline, [])

    def test_name_check(self):
        try:
            pipeline = PipelineConfigBuilder(None, "bad_dns")
            assert False
        except RuntimeError:
            pass
        except Exception:
            assert False

        pipeline2 = PipelineConfigBuilder(None, "good-dns")
        assert True
