# 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.
from __future__ import annotations

import json
from unittest import mock
from unittest.mock import AsyncMock, MagicMock, patch

import multidict
import pytest
import requests
from aiohttp import ClientResponseError, RequestInfo
from requests.exceptions import RequestException

from airflow.models import Connection
from airflow.providers.apache.livy.hooks.livy import BatchState, LivyAsyncHook, LivyHook
from airflow.providers.common.compat.sdk import AirflowException

from tests_common.test_utils.db import clear_test_connections

LIVY_CONN_ID = LivyHook.default_conn_name
DEFAULT_CONN_ID = LivyHook.default_conn_name
DEFAULT_HOST = "livy"
DEFAULT_SCHEMA = "http"
DEFAULT_PORT = 8998
MATCH_URL = f"//{DEFAULT_HOST}:{DEFAULT_PORT}"

BATCH_ID = 100
SAMPLE_GET_RESPONSE = {"id": BATCH_ID, "state": BatchState.SUCCESS.value}
VALID_SESSION_ID_TEST_CASES = [
    pytest.param(BATCH_ID, id="integer"),
    pytest.param(str(BATCH_ID), id="integer as string"),
]
INVALID_SESSION_ID_TEST_CASES = [
    pytest.param(None, id="none"),
    pytest.param("forty two", id="invalid string"),
    pytest.param({"a": "b"}, id="dictionary"),
]

INVALID_JAVA_SIZE_STRINGS = "Invalid java size format for string"
CONF_MUST_BE_DICT = "'conf' argument must be a dict"
CONF_VALUES_MUST_BE_STR_OR_INT = "'conf' values must be either strings or ints"


@pytest.mark.db_test
class TestLivyDbHook:
    @classmethod
    def setup_class(cls):
        clear_test_connections(add_default_connections_back=False)

    @classmethod
    def teardown_class(cls):
        clear_test_connections(add_default_connections_back=True)

    # TODO: Potential performance issue, converted setup_class to a setup_connections function level fixture
    @pytest.fixture(autouse=True)
    def setup_connections(self, create_connection_without_db):
        create_connection_without_db(
            Connection(
                conn_id=DEFAULT_CONN_ID,
                conn_type="http",
                host=DEFAULT_HOST,
                schema=DEFAULT_SCHEMA,
                port=DEFAULT_PORT,
            )
        )
        create_connection_without_db(Connection(conn_id="default_port", conn_type="http", host="http://host"))
        create_connection_without_db(Connection(conn_id="default_protocol", conn_type="http", host="host"))
        create_connection_without_db(Connection(conn_id="port_set", host="host", conn_type="http", port=1234))
        create_connection_without_db(
            Connection(conn_id="schema_set", host="host", conn_type="http", schema="https")
        )
        create_connection_without_db(
            Connection(conn_id="dont_override_schema", conn_type="http", host="http://host", schema="https")
        )
        create_connection_without_db(Connection(conn_id="missing_host", conn_type="http", port=1234))
        create_connection_without_db(Connection(conn_id="invalid_uri", uri="http://invalid_uri:4321"))
        create_connection_without_db(
            Connection(
                conn_id="with_credentials", login="login", password="secret", conn_type="http", host="host"
            )
        )

    @pytest.mark.parametrize(
        ("conn_id", "expected"),
        [
            pytest.param("default_port", "http://host", id="default-port"),
            pytest.param("default_protocol", "http://host", id="default-protocol"),
            pytest.param("port_set", "http://host:1234", id="with-defined-port"),
            pytest.param("schema_set", "https://host", id="with-defined-schema"),
            pytest.param("dont_override_schema", "http://host", id="ignore-defined-schema"),
        ],
    )
    def test_build_get_hook(self, conn_id, expected):
        hook = LivyHook(livy_conn_id=conn_id)
        hook.get_conn()
        assert hook.base_url == expected

    @pytest.mark.skip("Inherited HttpHook does not handle missing hostname")
    def test_missing_host(self):
        with pytest.raises(AirflowException):
            LivyHook(livy_conn_id="missing_host").get_conn()

    def test_build_body_minimal_request(self):
        assert LivyHook.build_post_batch_body(file="appname") == {"file": "appname"}

    def test_build_body_complex_request(self):
        body = LivyHook.build_post_batch_body(
            file="appname",
            class_name="org.example.livy",
            proxy_user="proxyUser",
            args=["a", "1"],
            jars=["jar1", "jar2"],
            files=["file1", "file2"],
            py_files=["py1", "py2"],
            archives=["arch1", "arch2"],
            queue="queue",
            name="name",
            conf={"a": "b"},
            driver_cores=2,
            driver_memory="1M",
            executor_memory="1m",
            executor_cores="1",
            num_executors="10",
        )

        assert body == {
            "file": "appname",
            "className": "org.example.livy",
            "proxyUser": "proxyUser",
            "args": ["a", "1"],
            "jars": ["jar1", "jar2"],
            "files": ["file1", "file2"],
            "pyFiles": ["py1", "py2"],
            "archives": ["arch1", "arch2"],
            "queue": "queue",
            "name": "name",
            "conf": {"a": "b"},
            "driverCores": 2,
            "driverMemory": "1M",
            "executorMemory": "1m",
            "executorCores": "1",
            "numExecutors": "10",
        }

    def test_parameters_validation(self):
        with pytest.raises(ValueError, match=INVALID_JAVA_SIZE_STRINGS):
            LivyHook.build_post_batch_body(file="appname", executor_memory="xxx")

        assert LivyHook.build_post_batch_body(file="appname", args=["a", 1, 0.1])["args"] == ["a", "1", "0.1"]

    @pytest.mark.parametrize(
        "size",
        [
            pytest.param("1m", id="lowercase-short"),
            pytest.param("1mb", id="lowercase-long"),
            pytest.param("1mb", id="uppercase-short"),
            pytest.param("1GB", id="uppercase-long"),
            pytest.param("1Gb", id="mix-case"),
            pytest.param(None, id="none"),
        ],
    )
    def test_validate_size_format(self, size):
        assert LivyHook._validate_size_format(size)

    @pytest.mark.parametrize(
        "size",
        [
            pytest.param("1Gb foo", id="fullmatch"),
            pytest.param("10", id="missing size"),
            pytest.param(1, id="integer"),
        ],
    )
    def test_validate_size_format_failed(self, size):
        with pytest.raises(ValueError, match=rf"Invalid java size format for string'{size}'"):
            assert LivyHook._validate_size_format(size)

    @pytest.mark.parametrize(
        "value",
        [
            pytest.param([1, "string"], id="list"),
            pytest.param((1, "string"), id="tuple"),
            pytest.param([], id="empty list"),
        ],
    )
    def test_validate_list_of_stringables(self, value):
        assert LivyHook._validate_list_of_stringables(value)

    @pytest.mark.parametrize(
        "value",
        [
            pytest.param({"a": "a"}, id="dict"),
            pytest.param([1, {}], id="invalid element"),
            pytest.param(None, id="none"),
            pytest.param(42, id="integer"),
            pytest.param("foo-bar", id="string"),
        ],
    )
    def test_validate_list_of_stringables_failed(self, value):
        with pytest.raises(ValueError, match="List of strings expected"):
            assert LivyHook._validate_list_of_stringables(value)

    @pytest.mark.parametrize(
        "config",
        [
            pytest.param({"k1": "v1", "k2": 0}, id="valid dictionary config"),
            pytest.param({}, id="empty dictionary"),
            pytest.param(None, id="none"),
        ],
    )
    def test_validate_extra_conf(self, config):
        LivyHook._validate_extra_conf(config)

    @pytest.mark.parametrize(
        "config",
        [
            pytest.param("k1=v1", id="string"),
            pytest.param([("k1", "v1"), ("k2", 0)], id="list of tuples"),
        ],
    )
    def test_validate_extra_conf_failed_non_dict(self, config):
        with pytest.raises(ValueError, match=CONF_MUST_BE_DICT):
            LivyHook._validate_extra_conf(config)

    @pytest.mark.parametrize(
        "config",
        [
            pytest.param({"outer": {"inner": "val"}}, id="nested dictionary"),
            pytest.param({"has_val": "val", "no_val": None}, id="none values in dictionary"),
            pytest.param({"has_val": "val", "no_val": ""}, id="empty values in dictionary"),
        ],
    )
    def test_validate_extra_conf_failed_dict(self, config):
        with pytest.raises(ValueError, match=CONF_VALUES_MUST_BE_STR_OR_INT):
            LivyHook._validate_extra_conf(config)

    @patch("airflow.providers.apache.livy.hooks.livy.LivyHook.run_method")
    def test_post_batch_arguments(self, mock_request):
        mock_request.return_value.status_code = 201
        mock_request.return_value.json.return_value = {
            "id": BATCH_ID,
            "state": BatchState.STARTING.value,
            "log": [],
        }

        resp = LivyHook().post_batch(file="sparkapp")

        mock_request.assert_called_once_with(
            method="POST", endpoint="/batches", data=json.dumps({"file": "sparkapp"}), headers={}
        )

        request_args = mock_request.call_args.kwargs
        assert "data" in request_args
        assert isinstance(request_args["data"], str)

        assert isinstance(resp, int)
        assert resp == BATCH_ID

    def test_post_batch_success(self, requests_mock):
        requests_mock.register_uri(
            "POST",
            "//livy:8998/batches",
            json={"id": BATCH_ID, "state": BatchState.STARTING.value, "log": []},
            status_code=201,
        )
        resp = LivyHook().post_batch(file="sparkapp")
        assert isinstance(resp, int)
        assert resp == BATCH_ID

    def test_post_batch_fail(self, requests_mock):
        requests_mock.register_uri("POST", f"{MATCH_URL}/batches", json={}, status_code=400, reason="ERROR")
        with pytest.raises(AirflowException):
            LivyHook().post_batch(file="sparkapp")

    def test_get_batch_success(self, requests_mock):
        requests_mock.register_uri(
            "GET", f"{MATCH_URL}/batches/{BATCH_ID}", json={"id": BATCH_ID}, status_code=200
        )
        resp = LivyHook().get_batch(BATCH_ID)
        assert isinstance(resp, dict)
        assert "id" in resp

    def test_get_batch_fail(self, requests_mock):
        requests_mock.register_uri(
            "GET",
            f"{MATCH_URL}/batches/{BATCH_ID}",
            json={"msg": "Unable to find batch"},
            status_code=404,
            reason="ERROR",
        )
        with pytest.raises(AirflowException):
            LivyHook().get_batch(BATCH_ID)

    @patch("airflow.providers.apache.livy.hooks.livy.LivyHook.run_method")
    @patch("airflow.providers.apache.livy.hooks.livy.LivyHook.get_conn")
    def test_post_batch_calls_get_conn_if_no_batch_id(self, mock_get_conn, mock_run_method):
        # mock run_method to get rid of call get_conn in it
        mock_response = MagicMock(resp=requests.Response)
        mock_response.json.return_value = {"id": BATCH_ID, "state": BatchState.STARTING.value, "log": []}
        mock_run_method.return_value = mock_response

        hook = LivyHook()

        # base_url is not set
        hook.post_batch(file="sparkapp")
        mock_get_conn.assert_called_once()

        # base_url is set
        mock_get_conn.reset_mock()
        hook.base_url = "//livy:8998"
        hook.post_batch(file="sparkapp")
        mock_get_conn.assert_not_called()

    def test_invalid_uri(self):
        with pytest.raises(RequestException):
            LivyHook(livy_conn_id="invalid_uri").post_batch(file="sparkapp")

    def test_get_batch_state_success(self, requests_mock):
        running = BatchState.RUNNING

        requests_mock.register_uri(
            "GET",
            f"{MATCH_URL}/batches/{BATCH_ID}/state",
            json={"id": BATCH_ID, "state": running.value},
            status_code=200,
        )

        state = LivyHook().get_batch_state(BATCH_ID)

        assert isinstance(state, BatchState)
        assert state == running

    def test_get_batch_state_fail(self, requests_mock):
        requests_mock.register_uri(
            "GET", f"{MATCH_URL}/batches/{BATCH_ID}/state", json={}, status_code=400, reason="ERROR"
        )
        with pytest.raises(AirflowException):
            LivyHook().get_batch_state(BATCH_ID)

    def test_get_batch_state_missing(self, requests_mock):
        requests_mock.register_uri("GET", f"{MATCH_URL}/batches/{BATCH_ID}/state", json={}, status_code=200)
        with pytest.raises(AirflowException):
            LivyHook().get_batch_state(BATCH_ID)

    def test_parse_post_response(self):
        res_id = LivyHook._parse_post_response({"id": BATCH_ID, "log": []})
        assert res_id == BATCH_ID

    def test_delete_batch_success(self, requests_mock):
        requests_mock.register_uri(
            "DELETE", f"{MATCH_URL}/batches/{BATCH_ID}", json={"msg": "deleted"}, status_code=200
        )
        assert LivyHook().delete_batch(BATCH_ID) == {"msg": "deleted"}

    def test_delete_batch_fail(self, requests_mock):
        requests_mock.register_uri(
            "DELETE", f"{MATCH_URL}/batches/{BATCH_ID}", json={}, status_code=400, reason="ERROR"
        )
        with pytest.raises(AirflowException):
            LivyHook().delete_batch(BATCH_ID)

    def test_missing_batch_id(self, requests_mock):
        requests_mock.register_uri("POST", f"{MATCH_URL}/batches", json={}, status_code=201)
        with pytest.raises(AirflowException):
            LivyHook().post_batch(file="sparkapp")

    @pytest.mark.parametrize("session_id", VALID_SESSION_ID_TEST_CASES)
    def test_get_batch_validation(self, session_id, requests_mock):
        requests_mock.register_uri(
            "GET", f"{MATCH_URL}/batches/{session_id}", json=SAMPLE_GET_RESPONSE, status_code=200
        )
        assert LivyHook().get_batch(session_id) == SAMPLE_GET_RESPONSE

    @pytest.mark.parametrize("session_id", INVALID_SESSION_ID_TEST_CASES)
    def test_get_batch_validation_failed(self, session_id):
        with pytest.raises(TypeError, match=r"\'session_id\' must be an integer"):
            LivyHook().get_batch(session_id)

    @pytest.mark.parametrize("session_id", VALID_SESSION_ID_TEST_CASES)
    def test_get_batch_state_validation(self, session_id, requests_mock):
        requests_mock.register_uri(
            "GET", f"{MATCH_URL}/batches/{session_id}/state", json=SAMPLE_GET_RESPONSE, status_code=200
        )
        assert LivyHook().get_batch_state(session_id) == BatchState.SUCCESS

    @pytest.mark.parametrize("session_id", INVALID_SESSION_ID_TEST_CASES)
    def test_get_batch_state_validation_failed(self, session_id):
        with pytest.raises(TypeError, match=r"\'session_id\' must be an integer"):
            LivyHook().get_batch_state(session_id)

    def test_delete_batch_validation(self, requests_mock):
        requests_mock.register_uri(
            "DELETE", f"{MATCH_URL}/batches/{BATCH_ID}", json={"id": BATCH_ID}, status_code=200
        )
        assert LivyHook().delete_batch(BATCH_ID) == {"id": BATCH_ID}

    @pytest.mark.parametrize("session_id", INVALID_SESSION_ID_TEST_CASES)
    def test_delete_batch_validation_failed(self, session_id):
        with pytest.raises(TypeError, match=r"\'session_id\' must be an integer"):
            LivyHook().delete_batch(session_id)

    @pytest.mark.parametrize("session_id", VALID_SESSION_ID_TEST_CASES)
    def test_check_session_id(self, session_id):
        LivyHook._validate_session_id(session_id)  # Should not raise any error

    @pytest.mark.parametrize("session_id", INVALID_SESSION_ID_TEST_CASES)
    def test_check_session_id_failed(self, session_id):
        with pytest.raises(TypeError, match=r"\'session_id\' must be an integer"):
            LivyHook._validate_session_id("asd")

    def test_extra_headers(self, requests_mock):
        requests_mock.register_uri(
            "POST",
            "//livy:8998/batches",
            json={"id": BATCH_ID, "state": BatchState.STARTING.value, "log": []},
            status_code=201,
            request_headers={"X-Requested-By": "user"},
        )

        hook = LivyHook(extra_headers={"X-Requested-By": "user"})
        hook.post_batch(file="sparkapp")

    def test_alternate_auth_type(self):
        auth_type = MagicMock()

        hook = LivyHook(livy_conn_id="with_credentials", auth_type=auth_type)

        auth_type.assert_not_called()

        hook.get_conn()

        auth_type.assert_called_once_with("login", "secret")

    @patch("airflow.providers.apache.livy.hooks.livy.LivyHook.run_method")
    def test_post_batch_with_endpoint_prefix(self, mock_request):
        mock_request.return_value.status_code = 201
        mock_request.return_value.json.return_value = {
            "id": BATCH_ID,
            "state": BatchState.STARTING.value,
            "log": [],
        }

        resp = LivyHook(endpoint_prefix="/livy").post_batch(file="sparkapp")

        mock_request.assert_called_once_with(
            method="POST", endpoint="/livy/batches", data=json.dumps({"file": "sparkapp"}), headers={}
        )

        request_args = mock_request.call_args.kwargs
        assert "data" in request_args
        assert isinstance(request_args["data"], str)

        assert isinstance(resp, int)
        assert resp == BATCH_ID

    def test_get_batch_with_endpoint_prefix(self, requests_mock):
        requests_mock.register_uri(
            "GET", f"{MATCH_URL}/livy/batches/{BATCH_ID}", json={"id": BATCH_ID}, status_code=200
        )
        resp = LivyHook(endpoint_prefix="/livy").get_batch(BATCH_ID)
        assert isinstance(resp, dict)
        assert "id" in resp

    def test_get_batch_state_with_endpoint_prefix(self, requests_mock):
        running = BatchState.RUNNING

        requests_mock.register_uri(
            "GET",
            f"{MATCH_URL}/livy/batches/{BATCH_ID}/state",
            json={"id": BATCH_ID, "state": running.value},
            status_code=200,
        )

        state = LivyHook(endpoint_prefix="/livy").get_batch_state(BATCH_ID)
        assert isinstance(state, BatchState)
        assert state == running

    def test_delete_batch_with_endpoint_prefix(self, requests_mock):
        requests_mock.register_uri(
            "DELETE", f"{MATCH_URL}/livy/batches/{BATCH_ID}", json={"msg": "deleted"}, status_code=200
        )
        assert LivyHook(endpoint_prefix="/livy").delete_batch(BATCH_ID) == {"msg": "deleted"}

    @pytest.mark.parametrize(
        "prefix",
        ["/livy/", "livy", "/livy", "livy/"],
        ids=["leading_and_trailing_slashes", "no_slashes", "leading_slash", "trailing_slash"],
    )
    def test_endpoint_prefix_is_sanitized_simple(self, requests_mock, prefix):
        requests_mock.register_uri(
            "GET", f"{MATCH_URL}/livy/batches/{BATCH_ID}", json={"id": BATCH_ID}, status_code=200
        )
        resp = LivyHook(endpoint_prefix=prefix).get_batch(BATCH_ID)
        assert isinstance(resp, dict)
        assert "id" in resp

    def test_endpoint_prefix_is_sanitized_multiple_path_elements(self, requests_mock):
        requests_mock.register_uri(
            "GET", f"{MATCH_URL}/livy/foo/bar/batches/{BATCH_ID}", json={"id": BATCH_ID}, status_code=200
        )
        resp = LivyHook(endpoint_prefix="/livy/foo/bar/").get_batch(BATCH_ID)
        assert isinstance(resp, dict)
        assert "id" in resp


class TestLivyAsyncHook:
    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_state_running(self, mock_run_method):
        """Asserts the batch state as running with success response."""
        mock_run_method.return_value = {"status": "success", "response": {"state": BatchState.RUNNING}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        state = await hook.get_batch_state(BATCH_ID)
        assert state == {
            "batch_state": BatchState.RUNNING,
            "response": "successfully fetched the batch state.",
            "status": "success",
        }

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_state_error(self, mock_run_method):
        """Asserts the batch state as error with error response."""
        mock_run_method.return_value = {"status": "error", "response": {"state": "error"}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        state = await hook.get_batch_state(BATCH_ID)
        assert state["status"] == "error"

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_state_error_without_state(self, mock_run_method):
        """Asserts the batch state as error without state returned as part of mock."""
        mock_run_method.return_value = {"status": "success", "response": {}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        state = await hook.get_batch_state(BATCH_ID)
        assert state["status"] == "error"

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_logs_success(self, mock_run_method):
        """Asserts the batch log as success."""
        mock_run_method.return_value = {"status": "success", "response": {}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        state = await hook.get_batch_logs(BATCH_ID, 0, 100)
        assert state["status"] == "success"

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_logs_error(self, mock_run_method):
        """Asserts the batch log for error."""
        mock_run_method.return_value = {"status": "error", "response": {}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        state = await hook.get_batch_logs(BATCH_ID, 0, 100)
        assert state["status"] == "error"

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.get_batch_logs")
    async def test_dump_batch_logs_success(self, mock_get_batch_logs):
        """Asserts the log dump log for success response."""
        mock_get_batch_logs.return_value = {
            "status": "success",
            "response": {"id": 1, "log": ["mock_log_1", "mock_log_2", "mock_log_3"]},
        }
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        log_dump = await hook.dump_batch_logs(BATCH_ID)
        assert log_dump == ["mock_log_1", "mock_log_2", "mock_log_3"]

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.get_batch_logs")
    async def test_dump_batch_logs_error(self, mock_get_batch_logs):
        """Asserts the log dump log for error response."""
        mock_get_batch_logs.return_value = {
            "status": "error",
            "response": {"id": 1, "log": ["mock_log_1", "mock_log_2"]},
        }
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        log_dump = await hook.dump_batch_logs(BATCH_ID)
        assert log_dump == {"id": 1, "log": ["mock_log_1", "mock_log_2"]}

    @pytest.mark.asyncio
    async def test_run_method_error(self):
        """Asserts the run_method for error response."""
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        response = await hook.run_method("localhost", "abc")
        assert response == {"status": "error", "response": "Invalid http method abc"}

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.http.hooks.http.aiohttp.ClientSession")
    @mock.patch(
        "airflow.providers.common.compat.connection.get_async_connection",
        return_value=Connection(
            conn_id=LIVY_CONN_ID,
            conn_type="http",
            host="http://host",
            port=80,
        ),
    )
    async def test_run_post_method_with_success(self, mock_get_connection, mock_session):
        """Asserts the run_method for success response for POST method."""
        mock_session.return_value.__aenter__.return_value.post = AsyncMock()
        mock_session.return_value.__aenter__.return_value.post.return_value.json = AsyncMock(
            return_value={"hello": "world"}
        )
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        response = await hook.run_method("api/jobs/runs/get")
        assert response["status"] == "success"
        assert response["response"] == {"hello": "world"}

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.http.hooks.http.aiohttp.ClientSession")
    @mock.patch(
        "airflow.providers.common.compat.connection.get_async_connection",
        return_value=Connection(
            conn_id=LIVY_CONN_ID,
            conn_type="http",
            host="http://host",
            port=80,
        ),
    )
    async def test_run_get_method_with_success(self, mock_get_connection, mock_session):
        """Asserts the run_method for GET method."""
        mock_session.return_value.__aenter__.return_value.get = AsyncMock()
        mock_session.return_value.__aenter__.return_value.get.return_value.json = AsyncMock(
            return_value={"hello": "world"}
        )
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        response = await hook.run_method("api/jobs/runs/get", method="GET")
        assert response["status"] == "success"
        assert response["response"] == {"hello": "world"}

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.http.hooks.http.aiohttp.ClientSession")
    @mock.patch(
        "airflow.providers.common.compat.connection.get_async_connection",
        return_value=Connection(
            conn_id=LIVY_CONN_ID,
            conn_type="http",
            host="http://host",
            port=80,
        ),
    )
    async def test_run_patch_method_with_success(self, mock_get_connection, mock_session):
        """Asserts the run_method for PATCH method."""
        mock_session.return_value.__aenter__.return_value.patch = AsyncMock()
        mock_session.return_value.__aenter__.return_value.patch.return_value.json = AsyncMock(
            return_value={"hello": "world"}
        )
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        response = await hook.run_method("api/jobs/runs/get", method="PATCH")
        assert response["status"] == "success"
        assert response["response"] == {"hello": "world"}

    @pytest.mark.asyncio
    @mock.patch(
        "airflow.providers.common.compat.connection.get_async_connection",
        return_value=Connection(
            conn_id=LIVY_CONN_ID,
            conn_type="http",
            host="http://host",
            port=80,
        ),
    )
    async def test_run_unexpected_method_with_success(self, mock_get_connection):
        """Asserts the run_method for unexpected method error"""
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        hook.method = "abc"
        response = await hook.run_method(endpoint="api/jobs/runs/get", headers={})
        assert response == {"response": "Invalid http method abc", "status": "error"}

    @pytest.mark.asyncio
    @mock.patch(
        "airflow.providers.common.compat.connection.get_async_connection",
        return_value=Connection(
            conn_id=LIVY_CONN_ID,
            conn_type="http",
            host="http://host",
            port=80,
        ),
    )
    async def test_run_put_method_with_type_error(self, mock_get_connection):
        """Asserts the run_method for TypeError."""

        async def mock_fun(arg1, arg2, arg3, arg4):
            return {"random value"}

        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID)
        hook.method = "PATCH"
        with pytest.raises(TypeError):
            await hook.run_method(endpoint="api/jobs/runs/get", data="test", headers=mock_fun)

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.http.hooks.http.aiohttp.ClientSession")
    @mock.patch(
        "airflow.providers.common.compat.connection.get_async_connection",
        return_value=Connection(
            conn_id=LIVY_CONN_ID,
            conn_type="http",
            host="http://host",
            port=80,
        ),
    )
    async def test_run_method_with_client_response_error(self, mock_get_connection, mock_session):
        """Asserts the run_method for Client Response Error."""

        mock_session.return_value.__aenter__.return_value.patch = AsyncMock(
            side_effect=ClientResponseError(
                request_info=RequestInfo(url="example.com", method="PATCH", headers=multidict.CIMultiDict()),
                status=500,
                history=[],
            )
        )
        hook = LivyAsyncHook(livy_conn_id="livy_default")
        hook.method = "PATCH"
        response = await hook.run_method("")
        assert response["status"] == "error"

    @pytest.fixture
    def setup_livy_conn(self, create_connection_without_db):
        create_connection_without_db(
            Connection(conn_id=LIVY_CONN_ID, conn_type="http", host="host", schema="http", port=8998)
        )
        create_connection_without_db(Connection(conn_id="default_port", conn_type="http", host="http://host"))
        create_connection_without_db(Connection(conn_id="default_protocol", conn_type="http", host="host"))
        create_connection_without_db(Connection(conn_id="port_set", host="host", conn_type="http", port=1234))
        create_connection_without_db(
            Connection(conn_id="schema_set", host="host", conn_type="http", schema="zzz")
        )
        create_connection_without_db(
            Connection(conn_id="dont_override_schema", conn_type="http", host="http://host", schema="zzz")
        )
        create_connection_without_db(Connection(conn_id="missing_host", conn_type="http", port=1234))
        create_connection_without_db(Connection(conn_id="invalid_uri", uri="http://invalid_uri:4321"))

    @pytest.mark.asyncio
    async def test_build_get_hook(self, setup_livy_conn):
        connection_url_mapping = {
            # id, expected
            "default_port": "http://host",
            "default_protocol": "http://host",
            "port_set": "http://host:1234",
            "schema_set": "zzz://host",
            "dont_override_schema": "http://host",
        }

        for conn_id, expected in connection_url_mapping.items():
            hook = LivyAsyncHook(livy_conn_id=conn_id)
            async with hook.session() as session:
                assert session.base_url == expected

    def test_build_body(self):
        # minimal request
        body = LivyAsyncHook.build_post_batch_body(file="appname")

        assert body == {"file": "appname"}

        # complex request
        body = LivyAsyncHook.build_post_batch_body(
            file="appname",
            class_name="org.example.livy",
            proxy_user="proxyUser",
            args=["a", "1"],
            jars=["jar1", "jar2"],
            files=["file1", "file2"],
            py_files=["py1", "py2"],
            archives=["arch1", "arch2"],
            queue="queue",
            name="name",
            conf={"a": "b"},
            driver_cores=2,
            driver_memory="1M",
            executor_memory="1m",
            executor_cores="1",
            num_executors="10",
        )

        assert body == {
            "file": "appname",
            "className": "org.example.livy",
            "proxyUser": "proxyUser",
            "args": ["a", "1"],
            "jars": ["jar1", "jar2"],
            "files": ["file1", "file2"],
            "pyFiles": ["py1", "py2"],
            "archives": ["arch1", "arch2"],
            "queue": "queue",
            "name": "name",
            "conf": {"a": "b"},
            "driverCores": 2,
            "driverMemory": "1M",
            "executorMemory": "1m",
            "executorCores": "1",
            "numExecutors": "10",
        }

    def test_parameters_validation(self):
        with pytest.raises(ValueError, match=INVALID_JAVA_SIZE_STRINGS):
            LivyAsyncHook.build_post_batch_body(file="appname", executor_memory="xxx")

        assert LivyAsyncHook.build_post_batch_body(file="appname", args=["a", 1, 0.1])["args"] == [
            "a",
            "1",
            "0.1",
        ]

    def test_parse_post_response(self):
        res_id = LivyAsyncHook._parse_post_response({"id": BATCH_ID, "log": []})

        assert res_id == BATCH_ID

    @pytest.mark.parametrize("valid_size", ["1m", "1mb", "1G", "1GB", "1Gb", None])
    def test_validate_size_format_success(self, valid_size):
        assert LivyAsyncHook._validate_size_format(valid_size)

    @pytest.mark.parametrize("invalid_size", ["1Gb foo", "10", 1])
    def test_validate_size_format_failure(self, invalid_size):
        with pytest.raises(ValueError, match=INVALID_JAVA_SIZE_STRINGS):
            assert LivyAsyncHook._validate_size_format(invalid_size)

    @pytest.mark.parametrize(
        "valid_string",
        [
            [1, "string"],
            (1, "string"),
            [],
        ],
    )
    def test_validate_list_of_stringables_success(self, valid_string):
        assert LivyAsyncHook._validate_list_of_stringables(valid_string)

    @pytest.mark.parametrize("invalid_string", [{"a": "a"}, [1, {}], [1, None], None, 1, "string"])
    def test_validate_list_of_stringables_failure(self, invalid_string):
        with pytest.raises(ValueError, match="List of strings expected"):
            LivyAsyncHook._validate_list_of_stringables(invalid_string)

    @pytest.mark.parametrize(
        "conf",
        [
            {"k1": "v1", "k2": 0},
            {},
            None,
        ],
    )
    def test_validate_extra_conf_success(self, conf):
        assert LivyAsyncHook._validate_extra_conf(conf)

    @pytest.mark.parametrize(
        "conf",
        [
            "k1=v1",
            [("k1", "v1"), ("k2", 0)],
        ],
    )
    def test_validate_extra_conf_failure_non_dict(self, conf):
        with pytest.raises(ValueError, match=CONF_MUST_BE_DICT):
            LivyAsyncHook._validate_extra_conf(conf)

    @pytest.mark.parametrize(
        "conf",
        [
            {"outer": {"inner": "val"}},
            {"has_val": "val", "no_val": None},
            {"has_val": "val", "no_val": ""},
        ],
    )
    def test_validate_extra_conf_failure_dict(self, conf):
        with pytest.raises(ValueError, match=CONF_VALUES_MUST_BE_STR_OR_INT):
            LivyAsyncHook._validate_extra_conf(conf)

    def test_parse_request_response(self):
        assert (
            LivyAsyncHook._parse_request_response(response={"id": BATCH_ID, "log": []}, parameter="id")
            == BATCH_ID
        )

    @pytest.mark.parametrize("conn_id", [100, 0])
    def test_check_session_id_success(self, conn_id):
        assert LivyAsyncHook._validate_session_id(conn_id) is None

    @pytest.mark.parametrize("conn_id", [None, "asd"])
    def test_check_session_id_failure(self, conn_id):
        with pytest.raises(TypeError):
            LivyAsyncHook._validate_session_id(None)

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_state_with_endpoint_prefix(self, mock_run_method):
        mock_run_method.return_value = {"status": "success", "response": {"state": BatchState.RUNNING}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID, endpoint_prefix="/livy")
        state = await hook.get_batch_state(BATCH_ID)
        assert state == {
            "batch_state": BatchState.RUNNING,
            "response": "successfully fetched the batch state.",
            "status": "success",
        }
        mock_run_method.assert_called_once_with(
            endpoint=f"/livy/batches/{BATCH_ID}/state",
            headers={},
        )

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_state_with_extra_headers(self, mock_run_method):
        headers = {"X-Requested-By": "user"}
        mock_run_method.return_value = {"status": "success", "response": {"state": BatchState.RUNNING}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID, extra_headers=headers)
        state = await hook.get_batch_state(BATCH_ID)
        assert state == {
            "batch_state": BatchState.RUNNING,
            "response": "successfully fetched the batch state.",
            "status": "success",
        }
        mock_run_method.assert_called_once_with(
            endpoint=f"/batches/{BATCH_ID}/state",
            headers=headers,
        )

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_logs_with_endpoint_prefix(self, mock_run_method):
        mock_run_method.return_value = {"status": "success", "response": {}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID, endpoint_prefix="/livy")
        state = await hook.get_batch_logs(BATCH_ID, 0, 100)
        assert state["status"] == "success"

        mock_run_method.assert_called_once_with(
            endpoint=f"/livy/batches/{BATCH_ID}/log",
            data={"from": 0, "size": 100},
            headers={},
        )

    @pytest.mark.asyncio
    @mock.patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.run_method")
    async def test_get_batch_logs_with_extra_headers(self, mock_run_method):
        headers = {"X-Requested-By": "user"}
        mock_run_method.return_value = {"status": "success", "response": {}}
        hook = LivyAsyncHook(livy_conn_id=LIVY_CONN_ID, extra_headers=headers)
        state = await hook.get_batch_logs(BATCH_ID, 0, 100)
        assert state["status"] == "success"

        mock_run_method.assert_called_once_with(
            endpoint=f"/batches/{BATCH_ID}/log",
            data={"from": 0, "size": 100},
            headers=headers,
        )
