#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
import json
import os
import unittest
from unittest import mock

from cloudpickle import cloudpickle

from ai_flow import WorkflowMeta
from ai_flow.frontend.web_server import generate_graph, AIFlowWebServerConfig, version
from ai_flow.test.scheduler_service.service.test_workflow_event_processor import MyContextExtractor
from ai_flow.version import __version__
from ai_flow.workflow.control_edge import WorkflowSchedulingRule, WorkflowAction, MeetAllEventCondition


class TestWebServer(unittest.TestCase):

    def test_version(self):
        self.assertEqual(version(), __version__)

    def test_generate_acyclic_graph(self):
        context_extractor = MyContextExtractor()
        rule = WorkflowSchedulingRule(MeetAllEventCondition().add_event('k', 'v', namespace='test_namespace'),
                                      WorkflowAction.STOP)
        workflow_meta = WorkflowMeta('workflow', 0, context_extractor_in_bytes=cloudpickle.dumps(context_extractor),
                                     scheduling_rules=[rule],
                                     graph='{"__af_object_type__": "jsonable", "__class__": "AIGraph", "__module__": "ai_flow.ai_graph.ai_graph", "_context_extractor": {"__af_object_type__": "jsonable", "__class__": "DailyWorkflowContextExtractor", "__module__": "__main__"}, "edges": {"AINode_0": [{"__af_object_type__": "jsonable", "__class__": "DataEdge", "__module__": "ai_flow.ai_graph.data_edge", "destination": "AINode_0", "port": 0, "source": "ReadDatasetNode_0"}], "AINode_1": [{"__af_object_type__": "jsonable", "__class__": "DataEdge", "__module__": "ai_flow.ai_graph.data_edge", "destination": "AINode_1", "port": 0, "source": "ReadDatasetNode_1"}], "AINode_2": [{"__af_object_type__": "jsonable", "__class__": "DataEdge", "__module__": "ai_flow.ai_graph.data_edge", "destination": "AINode_2", "port": 0, "source": "AINode_1"}], "WriteDatasetNode_1": [{"__af_object_type__": "jsonable", "__class__": "DataEdge", "__module__": "ai_flow.ai_graph.data_edge", "destination": "WriteDatasetNode_1", "port": 0, "source": "AINode_2"}], "daily_validate": [{"__af_object_type__": "jsonable", "__class__": "ControlEdge", "__module__": "ai_flow.workflow.control_edge", "destination": "daily_validate", "scheduling_rule": {"__af_object_type__": "jsonable", "__class__": "JobSchedulingRule", "__module__": "ai_flow.workflow.control_edge", "action": "START", "event_condition": {"__af_object_type__": "jsonable", "__class__": "MeetAnyEventCondition", "__module__": "ai_flow.workflow.control_edge", "condition_type": "MEET_ANY", "events": [{"__af_object_type__": "jsonable", "__class__": "EventMeetConfig", "__module__": "ai_flow.workflow.control_edge", "event_key": "daily_workflow.daily_training", "event_type": "JOB_STATUS_CHANGED", "event_value": "FINISHED", "life": "ONCE", "namespace": "workflow_on_event", "sender": "daily_training", "value_condition": "EQUALS"}]}}, "source": "*"}]}, "name": null, "node_id": "AIGraph_0", "nodes": {"AINode_0": {"__af_object_type__": "jsonable", "__class__": "AINode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "daily_training", "job_type": "python", "properties": {"entry_module_path": "daily_workflow"}}, "name": null, "node_config": {"base_model_info": null, "model_info": null, "name": null, "node_type": "train", "properties": null}, "node_id": "AINode_0", "output_num": 0, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003c__main__\nDailyTrainingTrain\nq\u0000)\u0081q\u0001.", "__module__": "builtins"}, "properties": {}}, "AINode_1": {"__af_object_type__": "jsonable", "__class__": "AINode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "daily_validate", "job_type": "python", "properties": {"entry_module_path": "daily_workflow"}}, "name": null, "node_config": {"name": null, "node_type": "transform", "properties": null}, "node_id": "AINode_1", "output_num": 1, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003c__main__\nDailyValidateTransformer\nq\u0000)\u0081q\u0001.", "__module__": "builtins"}, "properties": {}}, "AINode_2": {"__af_object_type__": "jsonable", "__class__": "AINode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "daily_validate", "job_type": "python", "properties": {"entry_module_path": "daily_workflow"}}, "name": null, "node_config": {"name": null, "node_type": "user_define_operation", "properties": null}, "node_id": "AINode_2", "output_num": 1, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003c__main__\nDailyValidate\nq\u0000)\u0081q\u0001}q\u0002(X\n\u0000\u0000\u0000model_nameq\u0003NX\n\u0000\u0000\u0000model_pathq\u0004NX\r\u0000\u0000\u0000model_versionq\u0005Nub.", "__module__": "builtins"}, "properties": {}}, "ReadDatasetNode_0": {"__af_object_type__": "jsonable", "__class__": "ReadDatasetNode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "daily_training", "job_type": "python", "properties": {"entry_module_path": "daily_workflow"}}, "name": null, "node_config": {"dataset": {"__af_object_type__": "jsonable", "__class__": "DatasetMeta", "__module__": "ai_flow.meta.dataset_meta", "catalog_connection_uri": null, "catalog_database": null, "catalog_name": null, "catalog_table": null, "catalog_type": null, "create_time": 1631259398354, "data_format": null, "description": null, "name": "daily_data", "properties": null, "schema": {"__af_object_type__": "jsonable", "__class__": "Schema", "__module__": "ai_flow.meta.dataset_meta", "name_list": null, "type_list": null}, "update_time": 1631259398354, "uri": "/tmp/daily_data", "uuid": 6}, "name": null, "node_type": "read_dataset", "properties": null}, "node_id": "ReadDatasetNode_0", "output_num": 1, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003c__main__\nDailyTrainingReader\nq\u0000)\u0081q\u0001.", "__module__": "builtins"}, "properties": {}}, "ReadDatasetNode_1": {"__af_object_type__": "jsonable", "__class__": "ReadDatasetNode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "daily_validate", "job_type": "python", "properties": {"entry_module_path": "daily_workflow"}}, "name": null, "node_config": {"dataset": {"__af_object_type__": "jsonable", "__class__": "DatasetMeta", "__module__": "ai_flow.meta.dataset_meta", "catalog_connection_uri": null, "catalog_database": null, "catalog_name": null, "catalog_table": null, "catalog_type": null, "create_time": 1631259398087, "data_format": null, "description": null, "name": "mnist_evaluate", "properties": null, "schema": {"__af_object_type__": "jsonable", "__class__": "Schema", "__module__": "ai_flow.meta.dataset_meta", "name_list": null, "type_list": null}, "update_time": 1631259398087, "uri": "/Users/sxnan/workspace/flinkml/flink-ai-extended/flink-ai-flow/examples/dataset_data/mnist_evaluate.npz", "uuid": 2}, "name": null, "node_type": "read_dataset", "properties": null}, "node_id": "ReadDatasetNode_1", "output_num": 1, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003c__main__\nDailyValidateReader\nq\u0000)\u0081q\u0001.", "__module__": "builtins"}, "properties": {}}, "WriteDatasetNode_0": {"__af_object_type__": "jsonable", "__class__": "WriteDatasetNode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "daily_training", "job_type": "python", "properties": {"entry_module_path": "daily_workflow"}}, "name": null, "node_config": {"dataset": {"__af_object_type__": "jsonable", "__class__": "DatasetMeta", "__module__": "ai_flow.meta.dataset_meta", "catalog_connection_uri": null, "catalog_database": null, "catalog_name": null, "catalog_table": null, "catalog_type": null, "create_time": 1631504656736, "data_format": null, "description": null, "name": "daily_train_result", "properties": null, "schema": {"__af_object_type__": "jsonable", "__class__": "Schema", "__module__": "ai_flow.meta.dataset_meta", "name_list": null, "type_list": null}, "update_time": 1631504656736, "uri": null, "uuid": 8}, "name": null, "node_type": "write_dataset", "properties": null}, "node_id": "WriteDatasetNode_0", "output_num": 0, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003c__main__\nDummyWriter\nq\u0000)\u0081q\u0001.", "__module__": "builtins"}, "properties": {}}, "WriteDatasetNode_1": {"__af_object_type__": "jsonable", "__class__": "WriteDatasetNode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "daily_validate", "job_type": "python", "properties": {"entry_module_path": "daily_workflow"}}, "name": null, "node_config": {"dataset": {"__af_object_type__": "jsonable", "__class__": "DatasetMeta", "__module__": "ai_flow.meta.dataset_meta", "catalog_connection_uri": null, "catalog_database": null, "catalog_name": null, "catalog_table": null, "catalog_type": null, "create_time": 1631504656808, "data_format": null, "description": null, "name": "daily_validate_result", "properties": null, "schema": {"__af_object_type__": "jsonable", "__class__": "Schema", "__module__": "ai_flow.meta.dataset_meta", "name_list": null, "type_list": null}, "update_time": 1631504656808, "uri": null, "uuid": 9}, "name": null, "node_type": "write_dataset", "properties": null}, "node_id": "WriteDatasetNode_1", "output_num": 0, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003c__main__\nDummyWriter\nq\u0000)\u0081q\u0001.", "__module__": "builtins"}, "properties": {}}}, "output_num": 0, "properties": {}}')
        workflow_graph = generate_graph(workflow_meta)
        self.assertIsNotNone(workflow_graph)
        graph_nodes = json.loads(workflow_graph)
        for graph_node in graph_nodes:
            if graph_node['id'] == 'daily_data':
                self.assertEqual(graph_node['layer'], 1)
            if graph_node['id'] == 'daily_train_result':
                self.assertEqual(graph_node['layer'], 1)
            if graph_node['id'] == 'daily_training':
                self.assertEqual(graph_node['layer'], 2)
            if graph_node['id'] == 'mnist_evaluate':
                self.assertEqual(graph_node['layer'], 2)
            if graph_node['id'] == 'daily_validate':
                self.assertEqual(graph_node['layer'], 3)
            if graph_node['id'] == 'daily_validate_result':
                self.assertEqual(graph_node['layer'], 4)

    def test_generate_ring_graph(self):
        context_extractor = MyContextExtractor()
        rule = WorkflowSchedulingRule(MeetAllEventCondition().add_event('k', 'v', namespace='test_namespace'),
                                      WorkflowAction.STOP)
        workflow_meta = WorkflowMeta('workflow', 0, context_extractor_in_bytes=cloudpickle.dumps(context_extractor),
                                     scheduling_rules=[rule],
                                     graph='{"__af_object_type__": "jsonable", "__class__": "AIGraph", "__module__": "ai_flow.ai_graph.ai_graph", "_context_extractor": {"__af_object_type__": "jsonable", "__class__": "BroadcastAllContextExtractor", "__module__": "ai_flow.api.context_extractor"}, "edges": {"task_2": [{"__af_object_type__": "jsonable", "__class__": "ControlEdge", "__module__": "ai_flow.workflow.control_edge", "destination": "task_2", "scheduling_rule": {"__af_object_type__": "jsonable", "__class__": "JobSchedulingRule", "__module__": "ai_flow.workflow.control_edge", "action": "START", "event_condition": {"__af_object_type__": "jsonable", "__class__": "MeetAnyEventCondition", "__module__": "ai_flow.workflow.control_edge", "condition_type": "MEET_ANY", "events": [{"__af_object_type__": "jsonable", "__class__": "EventMeetConfig", "__module__": "ai_flow.workflow.control_edge", "event_key": "simple_workflow.task_1", "event_type": "JOB_STATUS_CHANGED", "event_value": "FINISHED", "life": "ONCE", "namespace": "celery_examples", "sender": "task_1", "value_condition": "EQUALS"}]}}, "source": "*"}, {"__af_object_type__": "jsonable", "__class__": "ControlEdge", "__module__": "ai_flow.workflow.control_edge", "destination": "task_2", "scheduling_rule": {"__af_object_type__": "jsonable", "__class__": "JobSchedulingRule", "__module__": "ai_flow.workflow.control_edge", "action": "STOP", "event_condition": {"__af_object_type__": "jsonable", "__class__": "MeetAnyEventCondition", "__module__": "ai_flow.workflow.control_edge", "condition_type": "MEET_ANY", "events": [{"__af_object_type__": "jsonable", "__class__": "EventMeetConfig", "__module__": "ai_flow.workflow.control_edge", "event_key": "simple_workflow.task_3", "event_type": "JOB_STATUS_CHANGED", "event_value": "FINISHED", "life": "ONCE", "namespace": "celery_examples", "sender": "task_3", "value_condition": "EQUALS"}]}}, "source": "*"}], "task_3": [{"__af_object_type__": "jsonable", "__class__": "ControlEdge", "__module__": "ai_flow.workflow.control_edge", "destination": "task_3", "scheduling_rule": {"__af_object_type__": "jsonable", "__class__": "JobSchedulingRule", "__module__": "ai_flow.workflow.control_edge", "action": "START", "event_condition": {"__af_object_type__": "jsonable", "__class__": "MeetAnyEventCondition", "__module__": "ai_flow.workflow.control_edge", "condition_type": "MEET_ANY", "events": [{"__af_object_type__": "jsonable", "__class__": "EventMeetConfig", "__module__": "ai_flow.workflow.control_edge", "event_key": "simple_workflow.task_2", "event_type": "JOB_STATUS_CHANGED", "event_value": "RUNNING", "life": "ONCE", "namespace": "celery_examples", "sender": "task_2", "value_condition": "EQUALS"}]}}, "source": "*"}]}, "name": null, "node_id": "AIGraph_0", "nodes": {"AINode_0": {"__af_object_type__": "jsonable", "__class__": "AINode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "task_1", "job_type": "bash", "properties": {"entry_module_path": "simple_workflow"}}, "name": null, "node_config": {"name": null, "node_type": "user_define_operation", "properties": null}, "node_id": "AINode_0", "output_num": 1, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003cai_flow_plugins.job_plugins.bash.bash_processor\nBashProcessor\nq\u0000)\u0081q\u0001}q\u0002(X\f\u0000\u0000\u0000bash_commandq\u0003X\u0011\u0000\u0000\u0000echo before_sleepq\u0004X\u000f\u0000\u0000\u0000output_encodingq\u0005X\u0005\u0000\u0000\u0000utf-8q\u0006ub.", "__module__": "builtins"}, "properties": {}}, "AINode_1": {"__af_object_type__": "jsonable", "__class__": "AINode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "task_2", "job_type": "bash", "properties": {"entry_module_path": "simple_workflow"}}, "name": null, "node_config": {"name": null, "node_type": "user_define_operation", "properties": null}, "node_id": "AINode_1", "output_num": 1, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003cai_flow_plugins.job_plugins.bash.bash_processor\nBashProcessor\nq\u0000)\u0081q\u0001}q\u0002(X\f\u0000\u0000\u0000bash_commandq\u0003X\t\u0000\u0000\u0000sleep 100q\u0004X\u000f\u0000\u0000\u0000output_encodingq\u0005X\u0005\u0000\u0000\u0000utf-8q\u0006ub.", "__module__": "builtins"}, "properties": {}}, "AINode_2": {"__af_object_type__": "jsonable", "__class__": "AINode", "__module__": "ai_flow.ai_graph.ai_node", "config": {"__af_object_type__": "jsonable", "__class__": "JobConfig", "__module__": "ai_flow.workflow.job_config", "job_label_report_interval": 5.0, "job_name": "task_3", "job_type": "bash", "properties": {"entry_module_path": "simple_workflow"}}, "name": null, "node_config": {"name": null, "node_type": "user_define_operation", "properties": null}, "node_id": "AINode_2", "output_num": 1, "processor": {"__af_object_type__": "bytes", "__class__": "bytes", "__data__": "\u0080\u0003cai_flow_plugins.job_plugins.bash.bash_processor\nBashProcessor\nq\u0000)\u0081q\u0001}q\u0002(X\f\u0000\u0000\u0000bash_commandq\u0003X\b\u0000\u0000\u0000sleep 10q\u0004X\u000f\u0000\u0000\u0000output_encodingq\u0005X\u0005\u0000\u0000\u0000utf-8q\u0006ub.", "__module__": "builtins"}, "properties": {}}}, "output_num": 0, "properties": {}}')
        workflow_graph = generate_graph(workflow_meta)
        self.assertIsNotNone(workflow_graph)
        graph_nodes = json.loads(workflow_graph)
        for graph_node in graph_nodes:
            if graph_node['id'] == 'task_1':
                self.assertEqual(graph_node['layer'], 1)
            if graph_node['id'] == 'task_2':
                self.assertEqual(graph_node['layer'], 2)
            if graph_node['id'] == 'task_3':
                self.assertEqual(graph_node['layer'], 1)

    def test_web_server_read_config(self):
        import ai_flow.frontend.web_server as web_server
        aiflow_server_config_path = os.path.join(
            os.path.dirname(__file__), "test_aiflow_server.yaml"
        )

        with mock.patch('ai_flow.frontend.web_server.start_web_server') as start_web_server:
            web_server.main(['-c', aiflow_server_config_path])

            config = AIFlowWebServerConfig({'airflow_web_server_uri': 'http://localhost:8080', 'host': '127.0.0.1',
                                            'port': 8000})
            start_web_server.assert_called_once_with('scheduler_class_testing', 'db_uri_testing', config)


if __name__ == '__main__':
    unittest.main()
