# Copyright 2020 The StackStorm Authors.
# Copyright 2019 Extreme Networks, Inc.
#
# Licensed 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.

from __future__ import absolute_import
from unittest2 import TestCase
from mock import Mock

from st2common.transport.consumers import ActionsQueueConsumer
from st2common.models.db.liveaction import LiveActionDB

from st2tests import config as test_config

test_config.parse_args()

__all__ = ["ActionsQueueConsumerTestCase"]


class ActionsQueueConsumerTestCase(TestCase):
    def test_process_right_dispatcher_is_used(self):
        handler = Mock()
        handler.message_type = LiveActionDB
        consumer = ActionsQueueConsumer(connection=None, queues=None, handler=handler)

        # Non-workflow action
        consumer._workflows_dispatcher = Mock()
        consumer._actions_dispatcher = Mock()

        body = LiveActionDB(
            status="scheduled", action="core.local", action_is_workflow=False
        )
        message = Mock()
        consumer.process(body=body, message=message)

        self.assertEqual(consumer._workflows_dispatcher.dispatch.call_count, 0)
        self.assertEqual(consumer._actions_dispatcher.dispatch.call_count, 1)

        # Workflow action
        consumer._workflows_dispatcher = Mock()
        consumer._actions_dispatcher = Mock()

        body = LiveActionDB(
            status="scheduled", action="core.local", action_is_workflow=True
        )
        message = Mock()
        consumer.process(body=body, message=message)

        self.assertEqual(consumer._workflows_dispatcher.dispatch.call_count, 1)
        self.assertEqual(consumer._actions_dispatcher.dispatch.call_count, 0)

@@ -660,16 +660,98 @@
 _import%0A
+import random%0Aimport eventlet%0A%0Afrom kombu import Exchange%0Afrom kombu import Queue%0A
 from uni
@@ -859,16 +859,130 @@
 onsumer%0A
+from st2common.transport.publishers import PoolPublisher%0Afrom st2common.transport import utils as transport_utils%0A
 from st2
@@ -1191,16 +1191,71 @@
 tCase):%0A
+    message_count = 0%0A    message_type = LiveActionDB%0A%0A
     def 
@@ -2332,16 +2332,16 @@
 unt, 1)%0A
-
         
@@ -2390,28 +2390,1449 @@
 her.dispatch.call_count, 0)%0A
+%0A    def test_stop_consumption_on_shutdown(self):%0A        exchange = Exchange(%22st2.execution.test%22, type=%22topic%22)%0A        queue_name = %22test-%22 + str(random.randint(1, 10000))%0A        queue = Queue(%0A            name=queue_name, exchange=exchange, routing_key=%22#%22, auto_delete=True%0A        )%0A        publisher = PoolPublisher()%0A        with transport_utils.get_connection() as connection:%0A            connection.connect()%0A            watcher = ActionsQueueConsumer(%0A                connection=connection, queues=queue, handler=self%0A            )%0A            watcher_thread = eventlet.greenthread.spawn(watcher.run)%0A%0A        # Give it some time to start up since we are publishing on a new queue%0A        eventlet.sleep(0.5)%0A        body = LiveActionDB(%0A            status=%22scheduled%22, action=%22core.local%22, action_is_workflow=False%0A        )%0A        publisher.publish(payload=body, exchange=exchange)%0A        eventlet.sleep(0.2)%0A        self.assertEqual(self.message_count, 1)%0A        body = LiveActionDB(%0A            status=%22scheduled%22, action=%22core.local%22, action_is_workflow=True%0A        )%0A        watcher.shutdown()%0A        eventlet.sleep(1)%0A        publisher.publish(payload=body, exchange=exchange)%0A        # Second published message won't be consumed.%0A        self.assertEqual(self.message_count, 1)%0A        watcher_thread.kill()%0A%0A    def process(self, liveaction):%0A        self.message_count = self.message_count + 1%0A

Test case added