From 38ebd5fb20df6f2ce54c14709a25bab95f85da4f Mon Sep 17 00:00:00 2001
From: Dom Dwyer <dom@itsallbroken.com>
Date: Tue, 27 Sep 2022 10:59:06 +0200
Subject: [PATCH] test: simplify partition provider mock

Removes redundant fields from the MockPartitionProvider.
---
 ingester/src/data/partition.rs                | 12 ++++++++++++
 ingester/src/data/partition/resolver/mock.rs  | 14 +++-----------
 ingester/src/data/partition/resolver/trait.rs |  7 +------
 3 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/ingester/src/data/partition.rs b/ingester/src/data/partition.rs
index e232be2cdf..fad48c4d10 100644
--- a/ingester/src/data/partition.rs
+++ b/ingester/src/data/partition.rs
@@ -315,6 +315,18 @@ impl PartitionData {
     pub(crate) fn table_name(&self) -> &str {
         self.table_name.as_ref()
     }
+
+    /// Return the shard ID for this partition.
+    #[cfg(test)]
+    pub(crate) fn shard_id(&self) -> ShardId {
+        self.shard_id
+    }
+
+    /// Return the table ID for this partition.
+    #[cfg(test)]
+    pub(crate) fn table_id(&self) -> TableId {
+        self.table_id
+    }
 }
 
 #[cfg(test)]
diff --git a/ingester/src/data/partition/resolver/mock.rs b/ingester/src/data/partition/resolver/mock.rs
index c7e7a9dce3..f14e6850a9 100644
--- a/ingester/src/data/partition/resolver/mock.rs
+++ b/ingester/src/data/partition/resolver/mock.rs
@@ -23,27 +23,19 @@ impl MockPartitionProvider {
     pub(crate) fn with_partition(
         mut self,
         partition_key: PartitionKey,
-        shard_id: ShardId,
-        table_id: TableId,
         data: PartitionData,
     ) -> Self {
-        self.insert(partition_key, shard_id, table_id, data);
+        self.insert(partition_key, data);
         self
     }
 
     /// Add `data` to the mock state, returning it when asked for the specified
     /// `(key, shard, table)` triplet.
-    pub(crate) fn insert(
-        &mut self,
-        partition_key: PartitionKey,
-        shard_id: ShardId,
-        table_id: TableId,
-        data: PartitionData,
-    ) {
+    pub(crate) fn insert(&mut self, partition_key: PartitionKey, data: PartitionData) {
         assert!(
             self.partitions
                 .lock()
-                .insert((partition_key, shard_id, table_id), data)
+                .insert((partition_key, data.shard_id(), data.table_id()), data)
                 .is_none(),
             "overwriting an existing mock PartitionData"
         );
diff --git a/ingester/src/data/partition/resolver/trait.rs b/ingester/src/data/partition/resolver/trait.rs
index 19b2d99856..5d5d6df441 100644
--- a/ingester/src/data/partition/resolver/trait.rs
+++ b/ingester/src/data/partition/resolver/trait.rs
@@ -55,12 +55,7 @@ mod tests {
         let partition = PartitionId::new(4242);
         let data = PartitionData::new(partition, shard_id, table_id, Arc::clone(&table_name), None);
 
-        let mock = Arc::new(MockPartitionProvider::default().with_partition(
-            key.clone(),
-            shard_id,
-            table_id,
-            data,
-        ));
+        let mock = Arc::new(MockPartitionProvider::default().with_partition(key.clone(), data));
 
         let got = mock
             .get_partition(key, shard_id, table_id, Arc::clone(&table_name))