"""
Test Serverless Job Module
"""

import asyncio
from unittest.mock import Mock, patch

from unittest import IsolatedAsyncioTestCase
from aiohttp import ClientResponse
from aiohttp.test_utils import make_mocked_coro

from runpod.serverless.modules import rp_job


class TestJob(IsolatedAsyncioTestCase):
    """Tests the Job class."""

    async def test_get_job_200(self):
        """
        Tests the get_job function
        """
        # Mock the 200 response
        response4 = Mock()
        response4.status = 200
        response4.json = make_mocked_coro(
            return_value={"id": "123", "input": {"number": 1}}
        )

        with patch("aiohttp.ClientSession") as mock_session, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            # Set side_effect to a list of mock responses
            mock_session.get.return_value.__aenter__.side_effect = [response4]

            job = await rp_job.get_job(mock_session)

            # Assertions for the success case
            assert job == [{"id": "123", "input": {"number": 1}}]

    async def test_get_job_204(self):
        """
        Tests the get_job function with a 204 response
        """
        # 204 Mock
        response_204 = Mock()
        response_204.status = 204
        response_204.json = make_mocked_coro(return_value=None)

        with patch("aiohttp.ClientSession") as mock_session_204, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            mock_session_204.get.return_value.__aenter__.return_value = response_204
            job = await rp_job.get_job(mock_session_204)

            assert job is None
            assert mock_session_204.get.call_count == 1

    async def test_get_job_400(self):
        """
        Test the get_job function with a 400 response
        """
        # 400 Mock
        response_400 = Mock(ClientResponse)
        response_400.status = 400

        with patch("aiohttp.ClientSession") as mock_session_400, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            mock_session_400.get.return_value.__aenter__.return_value = response_400
            job = await rp_job.get_job(mock_session_400)

            assert job is None

    async def test_get_job_500(self):
        """
        Tests the get_job function with a 500 response
        """
        # 500 Mock
        response_500 = Mock(ClientResponse)
        response_500.status = 500

        with patch("aiohttp.ClientSession") as mock_session_500, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            mock_session_500.get.return_value.__aenter__.return_value = response_500
            job = await rp_job.get_job(mock_session_500)

            assert job is None

    async def test_get_job_no_id(self):
        """
        Tests the get_job function with a 200 response but no id
        """
        response = Mock(ClientResponse)
        response.status = 200
        response.json = make_mocked_coro(return_value={})

        with patch("aiohttp.ClientSession") as mock_session, patch(
            "runpod.serverless.modules.rp_job.log", new_callable=Mock
        ) as mock_log, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            mock_session.get.return_value.__aenter__.return_value = response

            job = await rp_job.get_job(mock_session)

            assert job == []
            assert mock_log.error.call_count == 1

    async def test_get_job_no_input(self):
        """
        Tests the get_job function with a 200 response but no input
        """
        response = Mock(ClientResponse)
        response.status = 200
        response.json = make_mocked_coro(return_value={"id": "123"})

        with patch("aiohttp.ClientSession") as mock_session, patch(
            "runpod.serverless.modules.rp_job.log", new_callable=Mock
        ) as mock_log, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            mock_session.get.return_value.__aenter__.return_value = response

            job = await rp_job.get_job(mock_session)

            assert job == []
            assert mock_log.error.call_count == 1

    async def test_get_job_no_timeout(self):
        """Tests the get_job function with a timeout"""
        # Timeout Mock
        response_timeout = Mock(ClientResponse)
        response_timeout.status = 200

        with patch("aiohttp.ClientSession") as mock_session_timeout, patch(
            "runpod.serverless.modules.rp_job.log", new_callable=Mock
        ) as mock_log, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            mock_session_timeout.get.return_value.__aenter__.side_effect = (
                asyncio.TimeoutError
            )
            job = await rp_job.get_job(mock_session_timeout)

            assert job == []
            assert mock_log.error.call_count == 0

    async def test_get_job_exception(self):
        """
        Tests the get_job function with an exception
        """
        # Exception Mock
        response_exception = Mock(ClientResponse)
        response_exception.status = 200

        with patch("aiohttp.ClientSession") as mock_session_exception, patch(
            "runpod.serverless.modules.rp_job.log", new_callable=Mock
        ) as mock_log, patch(
            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
        ):

            mock_session_exception.get.return_value.__aenter__.side_effect = Exception
            job = await rp_job.get_job(mock_session_exception)

            assert job == []
            assert mock_log.error.call_count == 1


class TestRunJob(IsolatedAsyncioTestCase):
    """Tests the run_job function"""

    def setUp(self) -> None:
        self.sample_job = {
            "id": "123",
            "input": {
                "test_input": None,
            },
        }

    async def test_simple_job(self):
        """
        Tests the run_job function
        """
        mock_handler = Mock()

        mock_handler.return_value = "test"
        job_result = await rp_job.run_job(mock_handler, self.sample_job)
        assert job_result == {"output": "test"}

        mock_handler.return_value = ["test1", "test2"]
        job_result_list = await rp_job.run_job(mock_handler, self.sample_job)
        assert job_result_list == {"output": ["test1", "test2"]}

        mock_handler.return_value = 123
        job_result_int = await rp_job.run_job(mock_handler, self.sample_job)
        assert job_result_int == {"output": 123}

    async def test_job_with_errors(self):
        """
        Tests the run_job function with errors
        """
        mock_handler = Mock()
        mock_handler.return_value = {"error": "test"}

        job_result = await rp_job.run_job(mock_handler, self.sample_job)

        assert job_result == {"error": "test"}

    async def test_job_with_raised_exception(self):
        """
        Tests the run_job function with a raised exception
        """
        mock_handler = Mock()
        mock_handler.side_effect = Exception

        job_result = await rp_job.run_job(mock_handler, self.sample_job)

        assert "error" in job_result

    async def test_job_with_refresh_worker(self):
        """
        Tests the run_job function with refresh_worker
        """
        mock_handler = Mock()
        mock_handler.return_value = {"refresh_worker": True}

        job_result = await rp_job.run_job(mock_handler, self.sample_job)

        assert job_result["stopPod"] is True

    async def test_job_bool_output(self):
        """
        Tests the run_job function with a boolean output
        """
        mock_handler = Mock()
        mock_handler.return_value = True

        job_result = await rp_job.run_job(mock_handler, self.sample_job)

        assert job_result == {"output": True}

    async def test_job_with_exception(self):
        """
        Tests the run_job function with an exception
        """
        mock_handler = Mock()
        mock_handler.side_effect = Exception

        job_result = await rp_job.run_job(mock_handler, self.sample_job)

        self.assertRaises(Exception, job_result)


class TestRunJobGenerator(IsolatedAsyncioTestCase):
    """Tests the run_job_generator function"""

    def handler_gen_success(self, job):  # pylint: disable=unused-argument
        """
        Test handler that returns a generator.
        """
        yield "partial_output_1"
        yield "partial_output_2"

    async def handler_async_gen_success(self, job):  # pylint: disable=unused-argument
        """
        Test handler that returns an async generator.
        """
        yield "partial_output_1"
        yield "partial_output_2"

    def handler_fail(self, job):
        """
        Test handler that raises an exception.
        """
        raise Exception("Test Exception")  # pylint: disable=broad-exception-raised

    async def test_run_job_generator_success(self):
        """
        Tests the run_job_generator function with a successful generator
        """
        handler = self.handler_gen_success
        job = {"id": "123"}

        with patch(
            "runpod.serverless.modules.rp_job.log", new_callable=Mock
        ) as mock_log:
            result = [i async for i in rp_job.run_job_generator(handler, job)]

        assert result == [
            {"output": "partial_output_1"},
            {"output": "partial_output_2"},
        ]
        assert mock_log.error.call_count == 0
        assert mock_log.info.call_count == 1
        mock_log.info.assert_called_with("Finished running generator.", "123")

    async def test_run_job_generator_success_async(self):
        """
        Tests the run_job_generator function with a successful generator
        """
        handler = self.handler_async_gen_success
        job = {"id": "123"}

        with patch(
            "runpod.serverless.modules.rp_job.log", new_callable=Mock
        ) as mock_log:
            result = [i async for i in rp_job.run_job_generator(handler, job)]

        assert result == [
            {"output": "partial_output_1"},
            {"output": "partial_output_2"},
        ]
        assert mock_log.error.call_count == 0
        assert mock_log.info.call_count == 1
        mock_log.info.assert_called_with("Finished running generator.", "123")

    async def test_run_job_generator_exception(self):
        """
        Tests the run_job_generator function with an exception
        """
        handler = self.handler_fail
        job = {"id": "123"}

        with patch(
            "runpod.serverless.modules.rp_job.log", new_callable=Mock
        ) as mock_log:
            result = [i async for i in rp_job.run_job_generator(handler, job)]

        assert len(result) == 1
        assert "error" in result[0]
        assert mock_log.error.call_count == 1
        assert mock_log.info.call_count == 1
        mock_log.info.assert_called_with("Finished running generator.", "123")
