import datetime
import tempfile
import unittest

import apache_beam as beam
import fastavro


class TestBeamSchemaConversions(unittest.TestCase):
    def test_convert_timestamp_millis(self):
        """Demonstrate bug: Avro-to-Beam schema conversion cannot handle timestamp-millis logical type."""
        avro_schema = {
            "type": "record",
            "name": "Test",
            "fields": [
                {"name": "name", "type": "string"},
                {"name": "timestamp", "type": {"type": "long", "logicalType": "timestamp-millis"}},
            ],
        }

        # Write test records into a temp file using fastavro.
        with tempfile.NamedTemporaryFile(delete=True) as input_avro_file:
            fastavro.writer(
                input_avro_file,
                avro_schema,
                [
                    {"name": "Alice", "timestamp": datetime.datetime(2024, 6, 3, 14, 14, 4, 765000, tzinfo=datetime.timezone.utc)},
                    {"name": "Bob", "timestamp": datetime.datetime(2024, 6, 3, 14, 14, 4, 765000, tzinfo=datetime.timezone.utc)},
                ],
                validator=True,
            )
            input_avro_file.flush()

            with self.assertRaises(TypeError) as context:
                with beam.Pipeline() as p:
                    avro_records = p | beam.io.ReadFromAvro(input_avro_file.name, as_rows=True)
                    avro_records | beam.LogElements()

            self.assertTrue("an integer is required" in str(context.exception))