Skip to main content

hegel/extras/jiff/
generators.rs

1use std::str::FromStr;
2
3use ciborium::Value;
4use jiff::civil::{Date, DateTime, Time};
5use jiff::tz::{Offset, TimeZone};
6use jiff::{SignedDuration, Span, Timestamp, Zoned};
7
8use crate::cbor_utils::{cbor_array, cbor_map};
9use crate::generators::{BasicGenerator, BoxedGenerator, Generator, deserialize_value};
10
11/// Generator for [`jiff::civil::Date`] values. Created by [`dates()`].
12pub struct DateGenerator;
13
14impl Generator<Date> for DateGenerator {
15    fn as_basic(&self) -> Option<BasicGenerator<'_, Date>> {
16        Some(BasicGenerator::new(cbor_map! {"type" => "date"}, |raw| {
17            let s: String = deserialize_value(raw);
18            Date::from_str(&s).unwrap()
19        }))
20    }
21}
22
23/// Generate [`jiff::civil::Date`] values.
24///
25/// See [`DateGenerator`] for builder methods.
26///
27/// # Example
28///
29/// ```no_run
30/// use hegel::extras::jiff as jiff_gs;
31///
32/// #[hegel::test]
33/// fn my_test(tc: hegel::TestCase) {
34///     let d = tc.draw(jiff_gs::dates());
35///     assert!(d.year() >= 1);
36/// }
37/// ```
38pub fn dates() -> DateGenerator {
39    DateGenerator
40}
41
42/// Generator for [`jiff::civil::Time`] values. Created by [`times()`].
43pub struct TimeGenerator;
44
45impl Generator<Time> for TimeGenerator {
46    fn as_basic(&self) -> Option<BasicGenerator<'_, Time>> {
47        Some(BasicGenerator::new(cbor_map! {"type" => "time"}, |raw| {
48            let s: String = deserialize_value(raw);
49            Time::from_str(&s).unwrap()
50        }))
51    }
52}
53
54/// Generate [`jiff::civil::Time`] values.
55///
56/// See [`TimeGenerator`] for builder methods.
57///
58/// # Example
59///
60/// ```no_run
61/// use hegel::extras::jiff as jiff_gs;
62///
63/// #[hegel::test]
64/// fn my_test(tc: hegel::TestCase) {
65///     let t = tc.draw(jiff_gs::times());
66///     assert!(t.hour() >= 0 && t.hour() <= 23);
67/// }
68/// ```
69pub fn times() -> TimeGenerator {
70    TimeGenerator
71}
72
73/// Convert a [`DateTime`] to nanoseconds since the Unix epoch, treating it
74/// as UTC. The wire format for [`DateTimeGenerator`] is a stable monotonic
75/// encoding of a civil datetime; UTC nanos give us that for free.
76///
77/// Bounds must fit within `Timestamp::MIN..=Timestamp::MAX`, which is slightly
78/// narrower than `DateTime::MIN..=DateTime::MAX` (jiff reserves a ~26-hour
79/// buffer at each end for offset conversions).
80fn datetime_to_nanos(dt: DateTime) -> i128 {
81    dt.to_zoned(TimeZone::UTC)
82        .expect("DateTime bound out of jiff::Timestamp range")
83        .timestamp()
84        .as_nanosecond()
85}
86
87/// Inverse of [`datetime_to_nanos`].
88fn nanos_to_datetime(n: i128) -> DateTime {
89    Timestamp::from_nanosecond(n)
90        .unwrap()
91        .to_zoned(TimeZone::UTC)
92        .datetime()
93}
94
95/// Generator for [`jiff::civil::DateTime`] values. Created by [`datetimes()`].
96pub struct DateTimeGenerator {
97    min_value: DateTime,
98    max_value: DateTime,
99}
100
101impl DateTimeGenerator {
102    /// Set the minimum datetime (inclusive).
103    pub fn min_value(mut self, min: DateTime) -> Self {
104        self.min_value = min;
105        self
106    }
107
108    /// Set the maximum datetime (inclusive).
109    pub fn max_value(mut self, max: DateTime) -> Self {
110        self.max_value = max;
111        self
112    }
113}
114
115impl Generator<DateTime> for DateTimeGenerator {
116    fn as_basic(&self) -> Option<BasicGenerator<'_, DateTime>> {
117        assert!(
118            self.min_value <= self.max_value,
119            "Cannot have max_value < min_value"
120        );
121        let schema = cbor_map! {
122            "type" => "integer",
123            "min_value" => datetime_to_nanos(self.min_value),
124            "max_value" => datetime_to_nanos(self.max_value),
125        };
126        Some(BasicGenerator::new(schema, |raw| {
127            let n: i128 = deserialize_value(raw);
128            nanos_to_datetime(n)
129        }))
130    }
131}
132
133/// Generate [`jiff::civil::DateTime`] values.
134///
135/// See [`DateTimeGenerator`] for builder methods.
136///
137/// # Example
138///
139/// ```no_run
140/// use hegel::extras::jiff as jiff_gs;
141/// use jiff::civil::DateTime;
142///
143/// #[hegel::test]
144/// fn my_test(tc: hegel::TestCase) {
145///     let min = DateTime::constant(2024, 1, 1, 0, 0, 0, 0);
146///     let dt = tc.draw(jiff_gs::datetimes().min_value(min));
147///     assert!(dt >= min);
148/// }
149/// ```
150pub fn datetimes() -> DateTimeGenerator {
151    DateTimeGenerator {
152        min_value: Timestamp::MIN.to_zoned(TimeZone::UTC).datetime(),
153        max_value: Timestamp::MAX.to_zoned(TimeZone::UTC).datetime(),
154    }
155}
156
157/// Generator for [`jiff::Timestamp`] values. Created by [`timestamps()`].
158pub struct TimestampGenerator {
159    min_value: Timestamp,
160    max_value: Timestamp,
161}
162
163impl TimestampGenerator {
164    /// Set the minimum timestamp (inclusive).
165    pub fn min_value(mut self, min: Timestamp) -> Self {
166        self.min_value = min;
167        self
168    }
169
170    /// Set the maximum timestamp (inclusive).
171    pub fn max_value(mut self, max: Timestamp) -> Self {
172        self.max_value = max;
173        self
174    }
175
176    fn build_schema(&self) -> Value {
177        assert!(
178            self.min_value <= self.max_value,
179            "Cannot have max_value < min_value"
180        );
181        cbor_map! {
182            "type" => "integer",
183            "min_value" => self.min_value.as_nanosecond(),
184            "max_value" => self.max_value.as_nanosecond(),
185        }
186    }
187}
188
189impl Generator<Timestamp> for TimestampGenerator {
190    fn as_basic(&self) -> Option<BasicGenerator<'_, Timestamp>> {
191        Some(BasicGenerator::new(self.build_schema(), |raw| {
192            let nanos: i128 = deserialize_value(raw);
193            Timestamp::from_nanosecond(nanos).unwrap()
194        }))
195    }
196}
197
198/// Generate [`jiff::Timestamp`] values.
199///
200/// See [`TimestampGenerator`] for builder methods.
201///
202/// # Example
203///
204/// ```no_run
205/// use hegel::extras::jiff as jiff_gs;
206/// use jiff::Timestamp;
207///
208/// #[hegel::test]
209/// fn my_test(tc: hegel::TestCase) {
210///     let t = tc.draw(jiff_gs::timestamps()
211///         .min_value(Timestamp::UNIX_EPOCH));
212///     assert!(t >= Timestamp::UNIX_EPOCH);
213/// }
214/// ```
215pub fn timestamps() -> TimestampGenerator {
216    TimestampGenerator {
217        min_value: Timestamp::MIN,
218        max_value: Timestamp::MAX,
219    }
220}
221
222/// The largest absolute nanosecond magnitude jiff allows in a `Span`.
223const SPAN_NANOS_LIMIT: i64 = i64::MAX;
224
225/// Generator for [`jiff::Span`] values. Created by [`spans()`].
226pub struct SpanGenerator {
227    min_nanos: i64,
228    max_nanos: i64,
229}
230
231impl SpanGenerator {
232    /// Set the minimum number of nanoseconds in the generated span.
233    pub fn min_nanoseconds(mut self, min: i64) -> Self {
234        self.min_nanos = min;
235        self
236    }
237
238    /// Set the maximum number of nanoseconds in the generated span.
239    pub fn max_nanoseconds(mut self, max: i64) -> Self {
240        self.max_nanos = max;
241        self
242    }
243
244    fn build_schema(&self) -> Value {
245        assert!(
246            self.min_nanos <= self.max_nanos,
247            "Cannot have max_nanoseconds < min_nanoseconds"
248        );
249        assert!(
250            self.min_nanos >= -SPAN_NANOS_LIMIT,
251            "min_nanoseconds must be >= -i64::MAX (the Span nanosecond limit)"
252        );
253        cbor_map! {
254            "type" => "integer",
255            "min_value" => self.min_nanos,
256            "max_value" => self.max_nanos,
257        }
258    }
259}
260
261impl Generator<Span> for SpanGenerator {
262    fn as_basic(&self) -> Option<BasicGenerator<'_, Span>> {
263        Some(BasicGenerator::new(self.build_schema(), |raw| {
264            let nanos: i64 = deserialize_value(raw);
265            Span::new().try_nanoseconds(nanos).unwrap()
266        }))
267    }
268}
269
270/// Generate [`jiff::Span`] values.
271///
272/// See [`SpanGenerator`] for builder methods.
273///
274/// # Example
275///
276/// ```no_run
277/// use hegel::extras::jiff as jiff_gs;
278///
279/// #[hegel::test]
280/// fn my_test(tc: hegel::TestCase) {
281///     let s = tc.draw(jiff_gs::spans().min_nanoseconds(0));
282///     assert!(s.get_nanoseconds() >= 0);
283/// }
284/// ```
285pub fn spans() -> SpanGenerator {
286    SpanGenerator {
287        min_nanos: -SPAN_NANOS_LIMIT,
288        max_nanos: SPAN_NANOS_LIMIT,
289    }
290}
291
292/// Inverse of [`SignedDuration::as_nanos`].
293fn nanos_to_signed_duration(n: i128) -> SignedDuration {
294    let secs = (n / 1_000_000_000) as i64;
295    let subnanos = (n % 1_000_000_000) as i32;
296    SignedDuration::new(secs, subnanos)
297}
298
299/// Generator for [`jiff::SignedDuration`] values. Created by [`signed_durations()`].
300pub struct SignedDurationGenerator {
301    min_value: SignedDuration,
302    max_value: SignedDuration,
303}
304
305impl SignedDurationGenerator {
306    /// Set the minimum duration (inclusive).
307    pub fn min_value(mut self, min: SignedDuration) -> Self {
308        self.min_value = min;
309        self
310    }
311
312    /// Set the maximum duration (inclusive).
313    pub fn max_value(mut self, max: SignedDuration) -> Self {
314        self.max_value = max;
315        self
316    }
317
318    fn build_schema(&self) -> Value {
319        assert!(
320            self.min_value <= self.max_value,
321            "Cannot have max_value < min_value"
322        );
323        cbor_map! {
324            "type" => "integer",
325            "min_value" => self.min_value.as_nanos(),
326            "max_value" => self.max_value.as_nanos(),
327        }
328    }
329}
330
331impl Generator<SignedDuration> for SignedDurationGenerator {
332    fn as_basic(&self) -> Option<BasicGenerator<'_, SignedDuration>> {
333        Some(BasicGenerator::new(self.build_schema(), |raw| {
334            let nanos: i128 = deserialize_value(raw);
335            nanos_to_signed_duration(nanos)
336        }))
337    }
338}
339
340/// Generate [`jiff::SignedDuration`] values.
341///
342/// # Example
343///
344/// ```no_run
345/// use hegel::extras::jiff as jiff_gs;
346/// use jiff::SignedDuration;
347///
348/// #[hegel::test]
349/// fn my_test(tc: hegel::TestCase) {
350///     let d = tc.draw(jiff_gs::signed_durations().min_value(SignedDuration::ZERO));
351///     assert!(d.as_nanos() >= 0);
352/// }
353/// ```
354pub fn signed_durations() -> SignedDurationGenerator {
355    SignedDurationGenerator {
356        min_value: SignedDuration::MIN,
357        max_value: SignedDuration::MAX,
358    }
359}
360
361/// The minimum representable offset in seconds (-25:59:59).
362const OFFSET_MIN_SECS: i32 = -93_599;
363/// The maximum representable offset in seconds (+25:59:59).
364const OFFSET_MAX_SECS: i32 = 93_599;
365
366/// Generator for [`jiff::tz::Offset`] values. Created by [`offsets()`].
367pub struct OffsetGenerator {
368    min_value: Offset,
369    max_value: Offset,
370}
371
372impl OffsetGenerator {
373    /// Set the minimum offset (inclusive).
374    pub fn min_value(mut self, min: Offset) -> Self {
375        self.min_value = min;
376        self
377    }
378
379    /// Set the maximum offset (inclusive).
380    pub fn max_value(mut self, max: Offset) -> Self {
381        self.max_value = max;
382        self
383    }
384
385    fn build_schema(&self) -> Value {
386        assert!(
387            self.min_value.seconds() <= self.max_value.seconds(),
388            "Cannot have max_value < min_value"
389        );
390        cbor_map! {
391            "type" => "integer",
392            "min_value" => self.min_value.seconds(),
393            "max_value" => self.max_value.seconds(),
394        }
395    }
396}
397
398impl Generator<Offset> for OffsetGenerator {
399    fn as_basic(&self) -> Option<BasicGenerator<'_, Offset>> {
400        Some(BasicGenerator::new(self.build_schema(), |raw| {
401            let secs: i32 = deserialize_value(raw);
402            Offset::from_seconds(secs).unwrap()
403        }))
404    }
405}
406
407/// Generate [`jiff::tz::Offset`] values.
408///
409/// See [`OffsetGenerator`] for builder methods.
410///
411/// # Example
412///
413/// ```no_run
414/// use hegel::extras::jiff as jiff_gs;
415/// use jiff::tz::Offset;
416///
417/// #[hegel::test]
418/// fn my_test(tc: hegel::TestCase) {
419///     let o = tc.draw(jiff_gs::offsets()
420///         .min_value(Offset::ZERO));
421///     assert!(o.seconds() >= 0);
422/// }
423/// ```
424pub fn offsets() -> OffsetGenerator {
425    OffsetGenerator {
426        min_value: Offset::from_seconds(OFFSET_MIN_SECS).unwrap(),
427        max_value: Offset::from_seconds(OFFSET_MAX_SECS).unwrap(),
428    }
429}
430
431/// Generator for [`jiff::Zoned`] values. Created by [`zoneds()`].
432pub struct ZonedGenerator<TS = TimestampGenerator, TZ = BoxedGenerator<'static, TimeZone>> {
433    timestamp_gen: TS,
434    timezone_gen: TZ,
435}
436
437impl<TS, TZ> ZonedGenerator<TS, TZ> {
438    /// Replace the timestamp generator.
439    pub fn timestamps<TS2>(self, timestamp_gen: TS2) -> ZonedGenerator<TS2, TZ>
440    where
441        TS2: Generator<Timestamp>,
442    {
443        ZonedGenerator {
444            timestamp_gen,
445            timezone_gen: self.timezone_gen,
446        }
447    }
448
449    /// Replace the timezone generator.
450    pub fn timezones<TZ2>(self, timezone_gen: TZ2) -> ZonedGenerator<TS, TZ2>
451    where
452        TZ2: Generator<TimeZone>,
453    {
454        ZonedGenerator {
455            timestamp_gen: self.timestamp_gen,
456            timezone_gen,
457        }
458    }
459}
460
461impl<TS, TZ> Generator<Zoned> for ZonedGenerator<TS, TZ>
462where
463    TS: Generator<Timestamp>,
464    TZ: Generator<TimeZone>,
465{
466    fn as_basic(&self) -> Option<BasicGenerator<'_, Zoned>> {
467        let ts_basic = self.timestamp_gen.as_basic()?;
468        let tz_basic = self.timezone_gen.as_basic()?;
469        let schema = cbor_map! {
470            "type" => "tuple",
471            "elements" => cbor_array![
472                ts_basic.schema().clone(),
473                tz_basic.schema().clone(),
474            ],
475        };
476        Some(BasicGenerator::new(schema, move |raw| {
477            let [ts_raw, tz_raw]: [Value; 2] = raw.into_array().unwrap().try_into().unwrap();
478            let ts = ts_basic.parse_raw(ts_raw);
479            let tz = tz_basic.parse_raw(tz_raw);
480            Zoned::new(ts, tz)
481        }))
482    }
483}
484
485/// Generate [`jiff::Zoned`] values.
486///
487/// # Example
488///
489/// ```no_run
490/// use hegel::extras::jiff as jiff_gs;
491/// use hegel::generators as gs;
492/// use jiff::tz::TimeZone;
493///
494/// #[hegel::test]
495/// fn my_test(tc: hegel::TestCase) {
496///     let _ = tc.draw(jiff_gs::zoneds());
497///     let _ = tc.draw(jiff_gs::zoneds().timezones(gs::just(TimeZone::UTC)));
498/// }
499/// ```
500pub fn zoneds() -> ZonedGenerator {
501    use crate::generators::DefaultGenerator;
502    ZonedGenerator {
503        timestamp_gen: timestamps(),
504        timezone_gen: <TimeZone as DefaultGenerator>::default_generator(),
505    }
506}