From 2dc31e738367f2013f6cae7e4397ab631d3e7565 Mon Sep 17 00:00:00 2001
From: morph <82043364+morph-dev@users.noreply.github.com>
Date: Thu, 15 Aug 2024 12:24:11 +0300
Subject: [PATCH] hivesim-rs: allow test data to be generic

---
 hivesim-rs/src/testapi.rs | 22 +++++++++++-----------
 hivesim-rs/src/types.rs   | 39 ---------------------------------------
 2 files changed, 11 insertions(+), 50 deletions(-)

diff --git a/hivesim-rs/src/testapi.rs b/hivesim-rs/src/testapi.rs
index 985632abb5..6324182d15 100644
--- a/hivesim-rs/src/testapi.rs
+++ b/hivesim-rs/src/testapi.rs
@@ -1,4 +1,4 @@
-use crate::types::{ClientDefinition, SuiteID, TestData, TestID, TestResult};
+use crate::types::{ClientDefinition, SuiteID, TestID, TestResult};
 use crate::Simulation;
 use ::std::{boxed::Box, future::Future, pin::Pin};
 use async_trait::async_trait;
@@ -21,9 +21,9 @@ pub type AsyncTestFunc = fn(
     >,
 >;
 
-pub type AsyncNClientsTestFunc = fn(
+pub type AsyncNClientsTestFunc<T> = fn(
     Vec<Client>,
-    Option<TestData>,
+    T,
 ) -> Pin<
     Box<
         dyn Future<Output = ()> // future API / pollable
@@ -202,7 +202,7 @@ pub async fn run_test(
 }
 
 #[derive(Clone)]
-pub struct NClientTestSpec {
+pub struct NClientTestSpec<T> {
     /// These fields are displayed in the UI. Be sure to add
     /// a meaningful description here.
     pub name: String,
@@ -212,17 +212,17 @@ pub struct NClientTestSpec {
     /// then perform further tests against it.
     pub always_run: bool,
     /// The Run function is invoked when the test executes.
-    pub run: AsyncNClientsTestFunc,
+    pub run: AsyncNClientsTestFunc<T>,
     /// For each client, there is a distinct map of Hive Environment Variable names to values.
     /// The environments must be in the same order as the `clients`
     pub environments: Option<Vec<Option<HashMap<String, String>>>>,
     /// test data which can be passed to the test
-    pub test_data: Option<TestData>,
+    pub test_data: T,
     pub clients: Vec<ClientDefinition>,
 }
 
 #[async_trait]
-impl Testable for NClientTestSpec {
+impl<T: Clone + Send + Sync + 'static> Testable for NClientTestSpec<T> {
     async fn run_test(&self, simulation: Simulation, suite_id: SuiteID, suite: Suite) {
         if let Some(test_match) = simulation.test_matcher.clone() {
             if !self.always_run && !test_match.match_test(&suite.name, &self.name) {
@@ -242,7 +242,7 @@ impl Testable for NClientTestSpec {
             simulation,
             test_run,
             self.environments.to_owned(),
-            self.test_data.to_owned(),
+            self.test_data.clone(),
             self.clients.to_owned(),
             self.run,
         )
@@ -251,13 +251,13 @@ impl Testable for NClientTestSpec {
 }
 
 // Write a test that runs against N clients.
-async fn run_n_client_test(
+async fn run_n_client_test<T: Send + 'static>(
     host: Simulation,
     test: TestRun,
     environments: Option<Vec<Option<HashMap<String, String>>>>,
-    test_data: Option<TestData>,
+    test_data: T,
     clients: Vec<ClientDefinition>,
-    func: AsyncNClientsTestFunc,
+    func: AsyncNClientsTestFunc<T>,
 ) {
     // Register test on simulation server and initialize the T.
     let test_id = host.start_test(test.suite_id, test.name, test.desc).await;
diff --git a/hivesim-rs/src/types.rs b/hivesim-rs/src/types.rs
index 66cf36e74b..c24e855084 100644
--- a/hivesim-rs/src/types.rs
+++ b/hivesim-rs/src/types.rs
@@ -36,42 +36,3 @@ pub struct TestResult {
     pub pass: bool,
     pub details: String,
 }
-
-#[derive(Clone, Debug)]
-pub struct ContentKeyValue {
-    pub key: String,
-    pub value: String,
-}
-
-#[derive(Clone, Debug)]
-pub struct ContentKeyOfferLookupValues {
-    pub key: String,
-    pub offer_value: String,
-    pub lookup_value: String,
-}
-
-#[derive(Clone, Debug)]
-pub enum TestData {
-    /// A list of tuples containing content key/value pairs
-    ContentList(Vec<ContentKeyValue>),
-    /// A list of tuples containing a content key, offer value, and lookup value
-    StateContentList(Vec<ContentKeyOfferLookupValues>),
-}
-
-impl TestData {
-    pub fn content_list(self) -> Vec<ContentKeyValue> {
-        if let TestData::ContentList(content_list) = self {
-            content_list
-        } else {
-            panic!("TestData didn't contain ContentList: enum was likely filled with the wrong data {self:?}")
-        }
-    }
-
-    pub fn state_content_list(self) -> Vec<ContentKeyOfferLookupValues> {
-        if let TestData::StateContentList(state_content_list) = self {
-            state_content_list
-        } else {
-            panic!("TestData didn't contain StateContentList: enum was likely filled with the wrong data {self:?}")
-        }
-    }
-}