"""Unit tests for SnowflakeMetadataConnector."""

from unittest import mock

import pytest

import snowflake_connector.metadata_connector as metadata_connector


@pytest.fixture(autouse=True)
def mock_requests_session():
    """Mock requests.Session."""
    with mock.patch(
            'snowflake_connector.metadata_connector.'
            'requests.Session') as session_mock:
        yield session_mock


def test_auth_by_password(sf_config_mock):
    """Test Snowflake authentication by password."""
    del sf_config_mock['private_key']

    connector = metadata_connector.SnowflakeMetadataConnector(sf_config_mock)
    connector.authenticate()

    connector.session.post.assert_called_once_with(
        'https://test_acc.snowflakecomputing.com/session/v1/login-request',
        json={
            'data': {
                'ACCOUNT_NAME': 'test_acc',
                'PASSWORD': 'test_pass',
                'LOGIN_NAME': 'test_user'}})


def test_auth_by_keypair(sf_config_mock, monkeypatch):
    """Test Snowflake authentication by password."""
    monkeypatch.setattr(
        metadata_connector.SnowflakeMetadataConnector, '_credentials',
        {'token': 'TEST_JWT_TOKEN'})
    connector = metadata_connector.SnowflakeMetadataConnector(sf_config_mock)
    connector.authenticate()

    connector.session.headers.update.assert_called_with(
        {'Authorization': 'Snowflake Token="TEST_JWT_TOKEN"'})


def test_cold_scan_bytes_number_fetch(sf_config_mock):
    """Test successful fetch of cold Snowflake query."""
    connector = metadata_connector.SnowflakeMetadataConnector(sf_config_mock)

    connector.session.get.return_value.json.return_value = {
        'data': {'queries': [{'stats': {'ioRemoteFdnReadBytes': 42}}]}}

    cursor = mock.MagicMock(sfqid='123-456')

    assert connector.get_query_scan_bytes_number(cursor) == 42
    connector.session.get.assert_called_once_with(
        'https://test_acc.snowflakecomputing.com/monitoring/queries/123-456')


def test_hot_scan_bytes_number_fetch(sf_config_mock):
    """Test successful fetch of hot Snowflake query."""
    connector = metadata_connector.SnowflakeMetadataConnector(sf_config_mock)

    connector.session.get.return_value.json.return_value = {
        'data': {'queries': [{'stats': {}}]}}

    cursor = mock.MagicMock(sfqid='123-456')

    assert connector.get_query_scan_bytes_number(cursor) == 0
    connector.session.get.assert_called_once_with(
        'https://test_acc.snowflakecomputing.com/monitoring/queries/123-456')


def test_bad_response_from_snowflake(sf_config_mock):
    """Test bad response from Snowflake raises SnowflakeBadResponse."""
    del sf_config_mock['private_key']
    connector = metadata_connector.SnowflakeMetadataConnector(sf_config_mock)

    connector.session.post.return_value.ok = False

    with pytest.raises(metadata_connector.SnowflakeBadResponse):
        connector.authenticate()