Repository: pydantic/pydantic-ai
File: tests/test_messages.py (Python)

Code before:
    assert video_url.media_type == 'video/mp4'
    assert video_url.format == 'mp4'


@pytest.mark.parametrize(
    'url,is_youtube',
    [
        pytest.param('https://youtu.be/lCdaVNyHtjU', True, id='youtu.be'),
        pytest.param('https://www.youtube.com/lCdaVNyHtjU', True, id='www.youtube.com'),
        pytest.param('https://youtube.com/lCdaVNyHtjU', True, id='youtube.com'),
        pytest.param('https://dummy.com/video.mp4', False, id='dummy.com'),
    ],
)
def test_youtube_video_url(url: str, is_youtube: bool):
    video_url = VideoUrl(url=url)
    assert video_url.is_youtube is is_youtube
    assert video_url.media_type == 'video/mp4'
    assert video_url.format == 'mp4'


def test_document_url():
    document_url = DocumentUrl(url='https://example.com/document.pdf')
    assert document_url.media_type == 'application/pdf'
    assert document_url.format == 'pdf'

    document_url = DocumentUrl(url='https://example.com/document', media_type_='application/pdf')
    assert document_url.media_type == 'application/pdf'
    assert document_url.format == 'pdf'


@pytest.mark.parametrize(
    'media_type, format',
    [
        ('audio/wav', 'wav'),
        ('audio/mpeg', 'mp3'),
    ],
)
def test_binary_content_audio(media_type: str, format: str):
    binary_content = BinaryContent(data=b'Hello, world!', media_type=media_type)
    assert binary_content.is_audio
    assert binary_content.format == format


@pytest.mark.parametrize(
    'media_type, format',
    [
        ('image/jpeg', 'jpeg'),
        ('image/png', 'png'),
        ('image/gif', 'gif'),
        ('image/webp', 'webp'),
    ],

Review comment (suggestion):
Can't we make the argument `media_type`? That may require a custom `__init__`.

Code after:
    assert video_url.media_type == 'video/mp4'
    assert video_url.format == 'mp4'


@pytest.mark.parametrize(
    'url,is_youtube',
    [
        pytest.param('https://youtu.be/lCdaVNyHtjU', True, id='youtu.be'),
        pytest.param('https://www.youtube.com/lCdaVNyHtjU', True, id='www.youtube.com'),
        pytest.param('https://youtube.com/lCdaVNyHtjU', True, id='youtube.com'),
        pytest.param('https://dummy.com/video.mp4', False, id='dummy.com'),
    ],
)
def test_youtube_video_url(url: str, is_youtube: bool):
    video_url = VideoUrl(url=url)
    assert video_url.is_youtube is is_youtube
    assert video_url.media_type == 'video/mp4'
    assert video_url.format == 'mp4'


def test_document_url():
    document_url = DocumentUrl(url='https://example.com/document.pdf')
    assert document_url.media_type == 'application/pdf'
    assert document_url.format == 'pdf'

    document_url = DocumentUrl(url='https://example.com/document', media_type='application/pdf')
    assert document_url.media_type == 'application/pdf'
    assert document_url.format == 'pdf'


@pytest.mark.parametrize(
    'media_type, format',
    [
        ('audio/wav', 'wav'),
        ('audio/mpeg', 'mp3'),
    ],
)
def test_binary_content_audio(media_type: str, format: str):
    binary_content = BinaryContent(data=b'Hello, world!', media_type=media_type)
    assert binary_content.is_audio
    assert binary_content.format == format


@pytest.mark.parametrize(
    'media_type, format',
    [
        ('image/jpeg', 'jpeg'),
        ('image/png', 'png'),
        ('image/gif', 'gif'),
        ('image/webp', 'webp'),
    ],