LangBot/tests/unit_tests/utils/test_image.py

"""
Unit tests for image utility functions.

Tests URL parsing and base64 extraction without network calls.
"""

from __future__ import annotations

import pytest
import base64

from langbot.pkg.utils.image import (
    decode_base64_limited,
    encode_base64,
    get_qq_image_downloadable_url,
    extract_b64_and_format,
)


@pytest.mark.asyncio
async def test_base64_media_helpers_round_trip_within_limit():
    encoded = await encode_base64(b'1234')

    assert await decode_base64_limited(encoded, max_bytes=4) == b'1234'


@pytest.mark.asyncio
async def test_base64_media_decode_rejects_oversized_payload():
    encoded = base64.b64encode(b'12345').decode()

    with pytest.raises(ValueError, match='exceeds'):
        await decode_base64_limited(encoded, max_bytes=4)


class TestGetQQImageDownloadableUrl:
    """Tests for get_qq_image_downloadable_url function."""

    def test_basic_url(self):
        """Parse basic image URL."""
        url = 'http://example.com/image.jpg'
        result_url, query = get_qq_image_downloadable_url(url)

        assert result_url == 'http://example.com/image.jpg'
        assert query == {}

    def test_url_with_query_params(self):
        """Parse URL with query parameters."""
        url = 'http://example.com/image.jpg?param1=value1&param2=value2'
        result_url, query = get_qq_image_downloadable_url(url)

        assert result_url == 'http://example.com/image.jpg'
        assert query == {'param1': ['value1'], 'param2': ['value2']}

    def test_url_with_port(self):
        """Parse URL with port number."""
        url = 'http://example.com:8080/image.jpg'
        result_url, query = get_qq_image_downloadable_url(url)

        assert result_url == 'http://example.com:8080/image.jpg'

    def test_url_with_path(self):
        """Parse URL with complex path."""
        url = 'http://example.com/path/to/image.jpg'
        result_url, query = get_qq_image_downloadable_url(url)

        assert result_url == 'http://example.com/path/to/image.jpg'

    def test_url_with_fragment(self):
        """Parse URL with fragment (fragment is not part of query)."""
        url = 'http://example.com/image.jpg#fragment'
        result_url, query = get_qq_image_downloadable_url(url)

        # Fragment is not included in query string parsing
        assert 'http://example.com/image.jpg' in result_url

    def test_https_url(self):
        """Parse HTTPS URL and preserve its scheme."""
        url = 'https://example.com/image.jpg'
        result_url, query = get_qq_image_downloadable_url(url)

        assert result_url == 'https://example.com/image.jpg'
        assert query == {}

    def test_preserves_qq_https_scheme_and_query(self):
        """QQ image URLs keep HTTPS and query parameters."""
        result_url, query = get_qq_image_downloadable_url('https://gchat.qpic.cn/gchatpic_new/abc/0?term=2&is_origin=1')

        assert result_url == 'https://gchat.qpic.cn/gchatpic_new/abc/0'
        assert query == {'term': ['2'], 'is_origin': ['1']}

    def test_defaults_missing_scheme_to_http(self):
        """Scheme-less image URLs default to HTTP."""
        result_url, query = get_qq_image_downloadable_url('gchat.qpic.cn/gchatpic_new/abc/0?term=2')

        assert result_url == 'http://gchat.qpic.cn/gchatpic_new/abc/0'
        assert query == {'term': ['2']}


class TestExtractB64AndFormat:
    """Tests for extract_b64_and_format function."""

    @pytest.mark.asyncio
    async def test_jpeg_data_uri(self):
        """Extract base64 and format from JPEG data URI."""
        # Create a simple base64 string
        original_data = b'test image data'
        b64_data = base64.b64encode(original_data).decode()
        data_uri = f'data:image/jpeg;base64,{b64_data}'

        result_b64, result_format = await extract_b64_and_format(data_uri)

        assert result_b64 == b64_data
        assert result_format == 'jpeg'

    @pytest.mark.asyncio
    async def test_png_data_uri(self):
        """Extract base64 and format from PNG data URI."""
        original_data = b'test png data'
        b64_data = base64.b64encode(original_data).decode()
        data_uri = f'data:image/png;base64,{b64_data}'

        result_b64, result_format = await extract_b64_and_format(data_uri)

        assert result_b64 == b64_data
        assert result_format == 'png'

    @pytest.mark.asyncio
    async def test_gif_data_uri(self):
        """Extract base64 and format from GIF data URI."""
        original_data = b'test gif data'
        b64_data = base64.b64encode(original_data).decode()
        data_uri = f'data:image/gif;base64,{b64_data}'

        result_b64, result_format = await extract_b64_and_format(data_uri)

        assert result_b64 == b64_data
        assert result_format == 'gif'

    @pytest.mark.asyncio
    async def test_webp_data_uri(self):
        """Extract base64 and format from WebP data URI."""
        original_data = b'test webp data'
        b64_data = base64.b64encode(original_data).decode()
        data_uri = f'data:image/webp;base64,{b64_data}'

        result_b64, result_format = await extract_b64_and_format(data_uri)

        assert result_b64 == b64_data
        assert result_format == 'webp'

    @pytest.mark.asyncio
    async def test_complex_base64(self):
        """Handle base64 with special characters."""
        # Base64 can include + and / characters
        original_data = bytes(range(256))  # All byte values
        b64_data = base64.b64encode(original_data).decode()
        data_uri = f'data:image/png;base64,{b64_data}'

        result_b64, result_format = await extract_b64_and_format(data_uri)

        assert result_b64 == b64_data
        # Verify we can decode back to original
        assert base64.b64decode(result_b64) == original_data

    @pytest.mark.asyncio
    async def test_empty_base64(self):
        """Handle empty base64 string."""
        data_uri = 'data:image/png;base64,'

        result_b64, result_format = await extract_b64_and_format(data_uri)

        assert result_b64 == ''
        assert result_format == 'png'