add dependency injection for timestamps

- fix all sorts of breakages in tests due to incorrect json assert - introduced some hacks, should be cleaned on this branch Signed-off-by: mimir-d <[email protected]>
opencomputeproject · Oct 8, 2024 · 358dde8 · 358dde8
1 parent 5f7dcf7
commit 358dde8
Show file tree

Hide file tree

Showing 6 changed files with 397 additions and 105 deletions.
diff --git a/src/output/config.rs b/src/output/config.rs
@@ -12,7 +12,7 @@ use crate::output::emitter;
 
 /// The configuration repository for the TestRun.
 pub struct Config {
-    pub(crate) timezone: chrono_tz::Tz,
+    pub(crate) timestamp_provider: Box<dyn TimestampProvider + Send + Sync + 'static>,
     pub(crate) writer: emitter::WriterType,
 }
 
@@ -32,20 +32,28 @@ impl Config {
 
 /// The builder for the [`Config`] object.
 pub struct ConfigBuilder {
-    timezone: Option<chrono_tz::Tz>,
+    timestamp_provider: Box<dyn TimestampProvider + Send + Sync + 'static>,
     writer: Option<emitter::WriterType>,
 }
 
 impl ConfigBuilder {
     fn new() -> Self {
         Self {
-            timezone: None,
+            timestamp_provider: Box::new(ConfiguredTzProvider { tz: chrono_tz::UTC }),
             writer: Some(emitter::WriterType::Stdout(emitter::StdoutWriter::new())),
         }
     }
 
     pub fn timezone(mut self, timezone: chrono_tz::Tz) -> Self {
-        self.timezone = Some(timezone);
+        self.timestamp_provider = Box::new(ConfiguredTzProvider { tz: timezone });
+        self
+    }
+
+    pub fn with_timestamp_provider(
+        mut self,
+        timestamp_provider: Box<dyn TimestampProvider + Send + Sync + 'static>,
+    ) -> Self {
+        self.timestamp_provider = timestamp_provider;
         self
     }
 
@@ -68,10 +76,37 @@ impl ConfigBuilder {
 
     pub fn build(self) -> Config {
         Config {
-            timezone: self.timezone.unwrap_or(chrono_tz::UTC),
+            timestamp_provider: self.timestamp_provider,
             writer: self
                 .writer
                 .unwrap_or(emitter::WriterType::Stdout(emitter::StdoutWriter::new())),
         }
     }
 }
+
+pub trait TimestampProvider {
+    fn now(&self) -> chrono::DateTime<chrono_tz::Tz>;
+
+    fn to_string(&self) -> String {
+        self.now()
+            .to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
+    }
+}
+
+struct ConfiguredTzProvider {
+    tz: chrono_tz::Tz,
+}
+
+impl TimestampProvider for ConfiguredTzProvider {
+    fn now(&self) -> chrono::DateTime<chrono_tz::Tz> {
+        chrono::Local::now().with_timezone(&self.tz)
+    }
+}
+
+pub struct NullTimestampProvider {}
+
+impl TimestampProvider for NullTimestampProvider {
+    fn now(&self) -> chrono::DateTime<chrono_tz::Tz> {
+        chrono::DateTime::from_timestamp_nanos(0).with_timezone(&chrono_tz::UTC)
+    }
+}
diff --git a/src/output/emitter.rs b/src/output/emitter.rs
@@ -16,6 +16,7 @@ use tokio::fs::File;
 use tokio::io::AsyncWriteExt;
 use tokio::sync::Mutex;
 
+use crate::output::config;
 use crate::spec;
 
 #[derive(Debug, thiserror::Error, derive_more::Display)]
@@ -84,26 +85,28 @@ impl StdoutWriter {
 }
 
 pub struct JsonEmitter {
-    timezone: chrono_tz::Tz,
+    // HACK: public for tests, but this should come from config directly to where needed
+    pub(crate) timestamp_provider: Box<dyn config::TimestampProvider + Send + Sync + 'static>,
     writer: WriterType,
     seqno: Arc<atomic::AtomicU64>,
 }
 
 impl JsonEmitter {
-    pub(crate) fn new(timezone: chrono_tz::Tz, writer: WriterType) -> Self {
+    pub(crate) fn new(
+        timestamp_provider: Box<dyn config::TimestampProvider + Send + Sync + 'static>,
+        writer: WriterType,
+    ) -> Self {
         JsonEmitter {
-            timezone,
+            timestamp_provider,
             writer,
             seqno: Arc::new(atomic::AtomicU64::new(0)),
         }
     }
 
     fn serialize_artifact(&self, object: &spec::RootImpl) -> serde_json::Value {
-        let now = chrono::Local::now();
-        let now_tz = now.with_timezone(&self.timezone);
         let root = spec::Root {
             artifact: object.clone(),
-            timestamp: now_tz,
+            timestamp: self.timestamp_provider.now(),
             seqno: self.incr_seqno(),
         };
         serde_json::json!(root)
@@ -144,7 +147,10 @@ mod tests {
 
         let buffer = Arc::new(Mutex::new(vec![]));
         let writer = BufferWriter::new(buffer.clone());
-        let emitter = JsonEmitter::new(chrono_tz::UTC, WriterType::Buffer(writer));
+        let emitter = JsonEmitter::new(
+            Box::new(config::NullTimestampProvider {}),
+            WriterType::Buffer(writer),
+        );
 
         emitter
             .emit(&spec::RootImpl::SchemaVersion(
@@ -179,7 +185,10 @@ mod tests {
 
         let buffer = Arc::new(Mutex::new(vec![]));
         let writer = BufferWriter::new(buffer.clone());
-        let emitter = JsonEmitter::new(chrono_tz::UTC, WriterType::Buffer(writer));
+        let emitter = JsonEmitter::new(
+            Box::new(config::NullTimestampProvider {}),
+            WriterType::Buffer(writer),
+        );
 
         let version = spec::RootImpl::SchemaVersion(spec::SchemaVersion::default());
         emitter.emit(&version).await?;

diff --git a/src/output/measure.rs b/src/output/measure.rs
@@ -180,7 +180,7 @@ impl StartedMeasurementSeries {
         let element = spec::MeasurementSeriesElement {
             index: self.incr_seqno(),
             value: value.clone(),
-            timestamp: chrono::Local::now().with_timezone(&chrono_tz::Tz::UTC),
+            timestamp: self.parent.emitter.timestamp_provider().now(),
             series_id: self.parent.start.series_id.clone(),
             metadata: None,
         };
@@ -223,7 +223,7 @@ impl StartedMeasurementSeries {
         let element = spec::MeasurementSeriesElement {
             index: self.incr_seqno(),
             value: value.clone(),
-            timestamp: chrono::Local::now().with_timezone(&chrono_tz::Tz::UTC),
+            timestamp: self.parent.emitter.timestamp_provider().now(),
             series_id: self.parent.start.series_id.clone(),
             metadata: Some(Map::from_iter(
                 metadata.iter().map(|(k, v)| (k.to_string(), v.clone())),

diff --git a/src/output/run.rs b/src/output/run.rs
@@ -258,7 +258,7 @@ impl TestRunBuilder {
 
     pub fn build(self) -> TestRun {
         let config = self.config.unwrap_or(config::Config::builder().build());
-        let emitter = emitter::JsonEmitter::new(config.timezone, config.writer);
+        let emitter = emitter::JsonEmitter::new(config.timestamp_provider, config.writer);
 
         TestRun {
             name: self.name,

diff --git a/src/output/step.rs b/src/output/step.rs
@@ -14,8 +14,8 @@ use crate::spec::{self, TestStepArtifactImpl};
 use tv::measure::MeasurementSeries;
 use tv::{emitter, error, log, measure};
 
-use super::JsonEmitter;
 use super::WriterError;
+use super::{JsonEmitter, TimestampProvider};
 
 /// A single test step in the scope of a [`TestRun`].
 ///
@@ -515,4 +515,8 @@ impl StepEmitter {
 
         Ok(())
     }
+
+    pub fn timestamp_provider(&self) -> &dyn TimestampProvider {
+        &*self.run_emitter.timestamp_provider
+    }
 }