From 353177a6339f044d18414a95215471757338d8c1 Mon Sep 17 00:00:00 2001
From: Paul Dicker <paul@pitdicker.nl>
Date: Fri, 1 Sep 2023 20:11:32 +0200
Subject: [PATCH] Add a feature and method to override `now()`

---
 Cargo.toml              |  3 ++-
 src/offset/local/mod.rs | 47 +++++++++++++++++++++++++++++++++++++++++
 src/offset/utc.rs       | 30 ++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 390b0e76ab..a17278561c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -27,6 +27,7 @@ oldtime = ["time"]
 wasmbind = ["wasm-bindgen", "js-sys"]
 unstable-locales = ["pure-rust-locales", "alloc"]
 __internal_bench = ["criterion"]
+test-override = []
 __doctest = []
 
 [dependencies]
@@ -66,7 +67,7 @@ doc-comment = { version = "0.3" }
 wasm-bindgen-test = "0.3"
 
 [package.metadata.docs.rs]
-features = ["serde"]
+features = ["serde", "test-override"]
 rustdoc-args = ["--cfg", "docsrs"]
 
 [package.metadata.playground]
diff --git a/src/offset/local/mod.rs b/src/offset/local/mod.rs
index 557f8db03e..91514b3a48 100644
--- a/src/offset/local/mod.rs
+++ b/src/offset/local/mod.rs
@@ -9,6 +9,8 @@ use rkyv::{Archive, Deserialize, Serialize};
 use super::fixed::FixedOffset;
 use super::{LocalResult, TimeZone};
 use crate::naive::{NaiveDate, NaiveDateTime, NaiveTime};
+#[cfg(feature = "test-override")]
+use crate::offset::utc::OVERRIDE_NOW;
 #[allow(deprecated)]
 use crate::Date;
 use crate::{DateTime, Utc};
@@ -145,8 +147,39 @@ impl Local {
     /// let now_with_offset = Local::now().with_timezone(&offset);
     /// ```
     pub fn now() -> DateTime<Local> {
+        #[cfg(all(feature = "test-override", test))]
+        if let Some(t) = OVERRIDE_NOW.with(|o| *o.borrow()) {
+            return t.into();
+        }
         Utc::now().with_timezone(&Local)
     }
+
+    /// Override the value that will be returned by `Local::now()` and `Utc::now(), for use in
+    /// tests.
+    ///
+    /// This method is only available behind the `test-override` feature, only works within the
+    /// current thread, and only in `cfg(test)`.
+    ///
+    /// ```no_run
+    /// # // Doctests don't have `cfg(test)`
+    /// use chrono::{Local, FixedOffset, TimeZone, Datelike};
+    /// fn is_today_leap_day() -> bool {
+    ///     let today = Local::now().date_naive();
+    /// dbg!(today);
+    ///     today.month() == 2 && today.day() == 29
+    /// }
+    ///
+    /// let now =
+    ///     FixedOffset::east_opt(3 * 60 * 60).unwrap().with_ymd_and_hms(2020, 2, 29, 0, 0, 0).unwrap();
+    /// Local::override_now(Some(now));
+    /// assert!(is_today_leap_day());
+    /// Local::override_now(None);
+    /// ```
+    #[cfg(feature = "test-override")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "test-override")))]
+    pub fn override_now(datetime: Option<DateTime<FixedOffset>>) {
+        OVERRIDE_NOW.with(|o| *o.borrow_mut() = datetime);
+    }
 }
 
 impl TimeZone for Local {
@@ -181,6 +214,8 @@ impl TimeZone for Local {
 mod tests {
     use super::Local;
     use crate::offset::TimeZone;
+    #[cfg(feature = "test-override")]
+    use crate::FixedOffset;
     use crate::{Datelike, Duration, Utc};
 
     #[test]
@@ -257,4 +292,16 @@ mod tests {
             );
         }
     }
+
+    #[cfg(all(feature = "clock", feature = "test-override"))]
+    #[test]
+    fn test_override() {
+        let now =
+            FixedOffset::east_opt(10800).unwrap().with_ymd_and_hms(2020, 2, 29, 12, 0, 0).unwrap();
+        Local::override_now(Some(now));
+        assert_eq!(Local::now(), now);
+
+        Local::override_now(None);
+        assert_ne!(Local::now(), now);
+    }
 }
diff --git a/src/offset/utc.rs b/src/offset/utc.rs
index d27f7877e1..bada8bc598 100644
--- a/src/offset/utc.rs
+++ b/src/offset/utc.rs
@@ -3,6 +3,8 @@
 
 //! The UTC (Coordinated Universal Time) time zone.
 
+#[cfg(all(feature = "clock", feature = "test-override"))]
+use core::cell::RefCell;
 use core::fmt;
 #[cfg(all(
     feature = "clock",
@@ -23,6 +25,12 @@ use crate::naive::{NaiveDate, NaiveDateTime};
 #[allow(deprecated)]
 use crate::{Date, DateTime};
 
+// Value to use for `Utc::now()` and `Local::now()`, when set with `Local::override_now`.
+#[cfg(all(feature = "clock", feature = "test-override"))]
+thread_local!(
+    pub(super) static OVERRIDE_NOW: RefCell<Option<DateTime<FixedOffset>>> = RefCell::new(None)
+);
+
 /// The UTC time zone. This is the most efficient time zone when you don't need the local time.
 /// It is also used as an offset (which is also a dummy type).
 ///
@@ -88,6 +96,10 @@ impl Utc {
     )))]
     #[must_use]
     pub fn now() -> DateTime<Utc> {
+        #[cfg(all(feature = "test-override", test))]
+        if let Some(t) = OVERRIDE_NOW.with(|o| *o.borrow()) {
+            return t.into();
+        }
         let now =
             SystemTime::now().duration_since(UNIX_EPOCH).expect("system time before Unix epoch");
         let naive =
@@ -147,3 +159,21 @@ impl fmt::Display for Utc {
         write!(f, "UTC")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[cfg(all(feature = "clock", feature = "test-override"))]
+    use crate::{FixedOffset, Local, TimeZone, Utc};
+
+    #[cfg(all(feature = "clock", feature = "test-override"))]
+    #[test]
+    fn test_override() {
+        let now =
+            FixedOffset::east_opt(10800).unwrap().with_ymd_and_hms(2020, 2, 29, 12, 0, 0).unwrap();
+        Local::override_now(Some(now));
+        assert_eq!(Utc::now(), now);
+
+        Local::override_now(None);
+        assert_ne!(Utc::now(), now);
+    }
+}