import os
import tempfile
from unittest.mock import MagicMock, call

import pytest

from airflow_tools.connectors.s3 import S3Connector


@pytest.fixture()
def test_folder():
    with tempfile.TemporaryDirectory() as temp_dir:
        file1 = os.path.join(temp_dir, "file1.txt")
        file2 = os.path.join(temp_dir, "file2.txt")

        # Create two empty files
        open(file1, "w").close()
        open(file2, "w").close()

        yield temp_dir


def test_upload_file():
    """Test the upload_file method of S3Connector."""
    mock_client = MagicMock()
    s3_connector = S3Connector()
    s3_connector._client = mock_client

    s3_connector.upload_file(
        bucket="test-bucket",
        filename="test-file.txt",
        key="test-key",
        kms_key_id="test-kms-key-id",
    )

    mock_client.upload_file.assert_called_once_with(
        Filename="test-file.txt",
        Bucket="test-bucket",
        Key="test-key",
        ExtraArgs={"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": "test-kms-key-id"},
    )


def test_upload_dir(test_folder):
    """Test the upload_dir method of S3Connector."""
    mock_client = MagicMock()
    s3_connector = S3Connector()
    s3_connector._client = mock_client

    s3_connector.upload_dir(
        bucket="test-bucket",
        folder=test_folder,
        key_prefix="test-prefix",
        kms_key_id="test-kms-key-id",
    )

    mock_client.upload_file.assert_has_calls(
        [
            call(
                Filename=f"{test_folder}/file1.txt",
                Bucket="test-bucket",
                Key="test-prefix/file1.txt",
                ExtraArgs={
                    "ServerSideEncryption": "aws:kms",
                    "SSEKMSKeyId": "test-kms-key-id",
                },
            ),
            call(
                Filename=f"{test_folder}/file2.txt",
                Bucket="test-bucket",
                Key="test-prefix/file2.txt",
                ExtraArgs={
                    "ServerSideEncryption": "aws:kms",
                    "SSEKMSKeyId": "test-kms-key-id",
                },
            ),
        ],
        any_order=True,
    )