From 9ac5d9f7f0c894892dfbdb9b20623464cb8018c1 Mon Sep 17 00:00:00 2001
From: Serial <69764315+Serial-ATA@users.noreply.github.com>
Date: Tue, 10 Dec 2024 14:28:05 -0500
Subject: [PATCH] chore: expand incredible-squaring e2e test

---
 blueprint-test-utils/src/test_ext.rs        |  4 +-
 blueprints/incredible-squaring/tests/e2e.rs | 82 ++++++++++++++++++---
 2 files changed, 74 insertions(+), 12 deletions(-)

diff --git a/blueprint-test-utils/src/test_ext.rs b/blueprint-test-utils/src/test_ext.rs
index 1ff17b65..c11f141d 100644
--- a/blueprint-test-utils/src/test_ext.rs
+++ b/blueprint-test-utils/src/test_ext.rs
@@ -183,7 +183,7 @@ pub async fn new_test_ext_blueprint_manager<
     match latest_revision {
         Some((rev, addr)) => debug!("MBSM is deployed at revision #{rev} at address {addr}"),
         None => {
-            let bytecode = MASTER_BLUEPRINT_SERVICE_MANAGER.to_vec();
+            let bytecode = MASTER_BLUEPRINT_SERVICE_MANAGER;
             gadget_sdk::trace!("Using MBSM bytecode of length: {}", bytecode.len());
 
             let ev = transactions::deploy_new_mbsm_revision(
@@ -191,7 +191,7 @@ pub async fn new_test_ext_blueprint_manager<
                 &client,
                 handles[0].sr25519_id(),
                 opts.signer_evm.clone().expect("Signer EVM is set"),
-                &bytecode,
+                bytecode,
             )
             .await
             .expect("deploy new MBSM revision");
diff --git a/blueprints/incredible-squaring/tests/e2e.rs b/blueprints/incredible-squaring/tests/e2e.rs
index d4ec359a..bf20e07b 100644
--- a/blueprints/incredible-squaring/tests/e2e.rs
+++ b/blueprints/incredible-squaring/tests/e2e.rs
@@ -1,14 +1,76 @@
-use blueprint_test_utils::{tangle::NodeConfig, test_tangle_blueprint, InputValue, OutputValue};
+use blueprint_test_utils::test_ext::new_test_ext_blueprint_manager;
+use blueprint_test_utils::{
+    tangle::NodeConfig, tempfile, wait_for_completion_of_tangle_job, InputValue, Job, OutputValue,
+};
 
 const SQUARING_JOB_ID: usize = 0;
 const N: usize = 5;
 
-test_tangle_blueprint!(
-    N, // Number of nodes
-    N,
-    SQUARING_JOB_ID,           // Job ID
-    [InputValue::Uint64(5)],   // Inputs
-    [OutputValue::Uint64(25)], // Expected output: input squared
-    0,
-    NodeConfig::new(false),
-);
+#[gadget_sdk::tokio::test(flavor = "multi_thread", crate = "::gadget_sdk::tokio")]
+async fn test_blueprint() {
+    blueprint_test_utils::setup_log();
+
+    let tmp_dir = tempfile::TempDir::new().unwrap();
+    let tmp_dir_path = tmp_dir.path().to_string_lossy().into_owned();
+
+    new_test_ext_blueprint_manager::<N, 1, String, _, _>(
+        tmp_dir_path,
+        blueprint_test_utils::run_test_blueprint_manager,
+        NodeConfig::new(false),
+    )
+    .await
+    .execute_with_async(|client, handles, blueprint, _| async move {
+        let keypair = handles[0].sr25519_id().clone();
+        let selected_service = &blueprint.services[0];
+        let service_id = selected_service.id;
+
+        gadget_sdk::info!("Submitting job {SQUARING_JOB_ID} with service ID {service_id}");
+
+        let job_args = vec![(InputValue::Uint64(5))];
+
+        let job = blueprint_test_utils::submit_job(
+            client,
+            &keypair,
+            service_id,
+            Job::from(SQUARING_JOB_ID as u8),
+            job_args,
+            0,
+        )
+        .await
+        .expect("Failed to submit job");
+
+        let call_id = job.call_id;
+
+        gadget_sdk::info!(
+            "Submitted job {SQUARING_JOB_ID} with service ID {service_id} has call id {call_id}"
+        );
+
+        let job_results = wait_for_completion_of_tangle_job(client, service_id, call_id, N)
+            .await
+            .expect("Failed to wait for job completion");
+
+        assert_eq!(job_results.service_id, service_id);
+        assert_eq!(job_results.call_id, call_id);
+
+        let expected_outputs = vec![OutputValue::Uint64(25)];
+        if expected_outputs.is_empty() {
+            gadget_sdk::info!("No expected outputs specified, skipping verification");
+            return;
+        }
+
+        assert_eq!(
+            job_results.result.len(),
+            expected_outputs.len(),
+            "Number of outputs doesn't match expected"
+        );
+
+        for (result, expected) in job_results
+            .result
+            .into_iter()
+            .zip(expected_outputs.into_iter())
+        {
+            assert_eq!(result, expected);
+        }
+    })
+    .await
+}