|
| 1 | +"""Tests for UTCDatetime pydantic type validation.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import importlib |
| 6 | +import inspect |
| 7 | +import pkgutil |
| 8 | +from datetime import UTC, datetime, timedelta, timezone |
| 9 | + |
| 10 | +import pytest |
| 11 | +from pydantic import BaseModel, ValidationError |
| 12 | + |
| 13 | +import diracx.core.models |
| 14 | +from diracx.core.models.types import UTCDatetime, _validate_utc |
| 15 | + |
| 16 | + |
| 17 | +class SampleModel(BaseModel): |
| 18 | + ts: UTCDatetime |
| 19 | + optional_ts: UTCDatetime | None = None |
| 20 | + |
| 21 | + |
| 22 | +class TestUTCDatetimeAcceptsUTC: |
| 23 | + def test_utc_timezone(self): |
| 24 | + dt = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC) |
| 25 | + m = SampleModel(ts=dt) |
| 26 | + assert m.ts == dt |
| 27 | + assert m.ts.tzinfo is UTC |
| 28 | + |
| 29 | + def test_timezone_utc(self): |
| 30 | + dt = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) |
| 31 | + m = SampleModel(ts=dt) |
| 32 | + assert m.ts.utcoffset() == timedelta(0) |
| 33 | + |
| 34 | + def test_iso_string_utc(self): |
| 35 | + m = SampleModel(ts="2024-01-01T12:00:00Z") |
| 36 | + assert m.ts.tzinfo is UTC |
| 37 | + |
| 38 | + def test_iso_string_plus_zero(self): |
| 39 | + m = SampleModel(ts="2024-01-01T12:00:00+00:00") |
| 40 | + assert m.ts.utcoffset() == timedelta(0) |
| 41 | + |
| 42 | + def test_optional_none(self): |
| 43 | + m = SampleModel(ts="2024-01-01T12:00:00Z", optional_ts=None) |
| 44 | + assert m.optional_ts is None |
| 45 | + |
| 46 | + |
| 47 | +class TestUTCDatetimeRejectsNonUTC: |
| 48 | + def test_naive_datetime(self): |
| 49 | + dt = datetime(2024, 1, 1, 12, 0, 0) # noqa: DTZ001 |
| 50 | + with pytest.raises(ValidationError, match="timezone"): |
| 51 | + SampleModel(ts=dt) |
| 52 | + |
| 53 | + def test_non_utc_timezone(self): |
| 54 | + cet = timezone(timedelta(hours=1)) |
| 55 | + dt = datetime(2024, 1, 1, 12, 0, 0, tzinfo=cet) |
| 56 | + with pytest.raises(ValidationError, match="must be in UTC"): |
| 57 | + SampleModel(ts=dt) |
| 58 | + |
| 59 | + def test_iso_string_non_utc(self): |
| 60 | + with pytest.raises(ValidationError, match="must be in UTC"): |
| 61 | + SampleModel(ts="2024-01-01T12:00:00+05:30") |
| 62 | + |
| 63 | + def test_naive_iso_string(self): |
| 64 | + with pytest.raises(ValidationError): |
| 65 | + SampleModel(ts="2024-01-01T12:00:00") |
| 66 | + |
| 67 | + |
| 68 | +def _is_datetime_type(annotation: type) -> bool: |
| 69 | + """Check if an annotation is datetime or a subclass of datetime.""" |
| 70 | + try: |
| 71 | + return isinstance(annotation, type) and issubclass(annotation, datetime) |
| 72 | + except TypeError: |
| 73 | + return False |
| 74 | + |
| 75 | + |
| 76 | +def _collect_model_classes() -> list[type[BaseModel]]: |
| 77 | + """Discover all BaseModel subclasses in diracx.core.models.""" |
| 78 | + models = [] |
| 79 | + package = diracx.core.models |
| 80 | + for _importer, modname, _ispkg in pkgutil.walk_packages( |
| 81 | + package.__path__, prefix=package.__name__ + "." |
| 82 | + ): |
| 83 | + if modname.endswith(".types"): |
| 84 | + continue |
| 85 | + module = importlib.import_module(modname) |
| 86 | + for _name, obj in inspect.getmembers(module, inspect.isclass): |
| 87 | + if ( |
| 88 | + issubclass(obj, BaseModel) |
| 89 | + and obj is not BaseModel |
| 90 | + and obj.__module__ == modname |
| 91 | + ): |
| 92 | + models.append(obj) |
| 93 | + return models |
| 94 | + |
| 95 | + |
| 96 | +def _check_field_uses_utc_validator(model: type[BaseModel], field_name: str) -> bool: |
| 97 | + """Check that a datetime field has the _validate_utc AfterValidator.""" |
| 98 | + field_info = model.model_fields[field_name] |
| 99 | + return any(getattr(m, "func", None) is _validate_utc for m in field_info.metadata) |
| 100 | + |
| 101 | + |
| 102 | +def test_all_datetime_fields_use_utc_datetime(): |
| 103 | + """Ensure no pydantic model in diracx.core.models uses bare datetime. |
| 104 | +
|
| 105 | + Every datetime field must use UTCDatetime to enforce UTC validation. |
| 106 | + """ |
| 107 | + violations = [] |
| 108 | + for model in _collect_model_classes(): |
| 109 | + for field_name, field_info in model.model_fields.items(): |
| 110 | + if not _is_datetime_type(field_info.annotation): |
| 111 | + continue |
| 112 | + if not _check_field_uses_utc_validator(model, field_name): |
| 113 | + violations.append(f"{model.__name__}.{field_name}") |
| 114 | + |
| 115 | + assert not violations, ( |
| 116 | + "The following fields use bare datetime instead of UTCDatetime:\n" |
| 117 | + + "\n".join(f" - {v}" for v in violations) |
| 118 | + ) |
0 commit comments