test: MockDmlHandler generic over write input

Allow the MockDmlHandler to capture any input type given to the write()
method. This lets us reuse the mock across all handler implementations,
regardless of their expected write input type.

router2/src/dml_handlers/mock.rs

@@ -1,20 +1,20 @@
-use std::{collections::VecDeque, sync::Arc};
+use std::{collections::VecDeque, fmt::Debug, sync::Arc};
 
 use async_trait::async_trait;
 use data_types::{delete_predicate::DeletePredicate, DatabaseName};
 
-use hashbrown::HashMap;
-use mutable_batch::MutableBatch;
 use parking_lot::Mutex;
 use trace::ctx::SpanContext;
 
 use super::{DmlError, DmlHandler};
 
+/// A captured call to a [`MockDmlHandler`], generic over `W`, the captured
+/// [`DmlHandler::WriteInput`] type.
 #[derive(Debug, Clone)]
-pub enum MockDmlHandlerCall {
+pub enum MockDmlHandlerCall<W> {
     Write {
         namespace: String,
-        batches: HashMap<String, MutableBatch>,
+        write_input: W,
     },
     Delete {
         namespace: String,
@@ -23,23 +23,42 @@ pub enum MockDmlHandlerCall {
     },
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
-struct Inner {
+struct Inner<W> {
-    calls: Vec<MockDmlHandlerCall>,
+    calls: Vec<MockDmlHandlerCall<W>>,
     write_return: VecDeque<Result<(), DmlError>>,
     delete_return: VecDeque<Result<(), DmlError>>,
 }
 
-impl Inner {
+impl<W> Default for Inner<W> {
-    fn record_call(&mut self, call: MockDmlHandlerCall) {
+    fn default() -> Self {
+        Self {
+            calls: Default::default(),
+            write_return: Default::default(),
+            delete_return: Default::default(),
+        }
+    }
+}
+
+impl<W> Inner<W> {
+    fn record_call(&mut self, call: MockDmlHandlerCall<W>) {
         self.calls.push(call);
     }
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
-pub struct MockDmlHandler(Mutex<Inner>);
+pub struct MockDmlHandler<W>(Mutex<Inner<W>>);
 
-impl MockDmlHandler {
+impl<W> Default for MockDmlHandler<W> {
+    fn default() -> Self {
+        Self(Default::default())
+    }
+}
+
+impl<W> MockDmlHandler<W>
+where
+    W: Clone,
+{
     pub fn with_write_return(self, ret: impl Into<VecDeque<Result<(), DmlError>>>) -> Self {
         self.0.lock().write_return = ret.into();
         self
@@ -50,7 +69,7 @@ impl MockDmlHandler {
         self
     }
 
-    pub fn calls(&self) -> Vec<MockDmlHandlerCall> {
+    pub fn calls(&self) -> Vec<MockDmlHandlerCall<W>> {
         self.0.lock().calls.clone()
     }
 }
@@ -68,22 +87,25 @@ macro_rules! record_and_return {
 }
 
 #[async_trait]
-impl DmlHandler for Arc<MockDmlHandler> {
+impl<W> DmlHandler for Arc<MockDmlHandler<W>>
+where
+    W: Debug + Send + Sync,
+{
     type WriteError = DmlError;
     type DeleteError = DmlError;
-    type WriteInput = HashMap<String, MutableBatch>;
+    type WriteInput = W;
 
     async fn write(
         &self,
         namespace: DatabaseName<'static>,
-        batches: Self::WriteInput,
+        write_input: Self::WriteInput,
         _span_ctx: Option<SpanContext>,
     ) -> Result<(), Self::WriteError> {
         record_and_return!(
             self,
             MockDmlHandlerCall::Write {
                 namespace: namespace.into(),
-                batches,
+                write_input,
             },
             write_return
         )

router2/src/dml_handlers/schema_validation.rs

@@ -363,9 +363,9 @@ mod tests {
         assert_matches!(err, SchemaError::Validate(_));
 
         // THe mock should observe exactly one write from the first call.
-        assert_matches!(mock.calls().as_slice(), [MockDmlHandlerCall::Write{namespace, batches}] => {
+        assert_matches!(mock.calls().as_slice(), [MockDmlHandlerCall::Write{namespace, write_input}] => {
             assert_eq!(namespace, NAMESPACE);
-            let batch = batches.get("bananas").expect("table not found in write");
+            let batch = write_input.get("bananas").expect("table not found in write");
             assert_eq!(batch.rows(), 1);
             let col = batch.column("val").expect("column not found in write");
             assert_matches!(col.influx_type(), InfluxColumnType::Field(InfluxFieldType::Integer));