From bf692a1417345e78d95369785f1ca2ffa6ddee22 Mon Sep 17 00:00:00 2001
From: Shanin Roman <shanin1000@yandex.ru>
Date: Tue, 4 Jun 2024 10:52:07 +0300
Subject: [PATCH] fix(client): allow to request QueryBox

Signed-off-by: Shanin Roman <shanin1000@yandex.ru>
---
 client/src/client.rs         | 21 ++++++++++++++++++++-
 core/test_network/src/lib.rs | 12 +++++++-----
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/client/src/client.rs b/client/src/client.rs
index d132c67234a..013e30403b7 100644
--- a/client/src/client.rs
+++ b/client/src/client.rs
@@ -229,7 +229,7 @@ impl From<ResponseReport> for eyre::Report {
 /// Output of a query
 pub trait QueryOutput: Into<QueryOutputBox> + TryFrom<QueryOutputBox> {
     /// Type of the query output
-    type Target: Clone;
+    type Target;
 
     /// Construct query output from query response
     fn new(output: Self, query_request: QueryResponseHandler<Self>) -> Self::Target;
@@ -309,6 +309,25 @@ where
     }
 }
 
+impl QueryOutput for QueryOutputBox {
+    type Target = QueryResult<Self>;
+
+    fn new(output: Self, query_handler: QueryResponseHandler<Self>) -> Self::Target {
+        match output {
+            Self::Vec(v) => {
+                let query_handler = QueryResponseHandler::new(query_handler.query_request);
+                let set = ResultSet {
+                    query_handler,
+                    iter: v,
+                    client_cursor: 0,
+                };
+                set.collect::<QueryResult<Vec<_>>>().map(Self::Vec)
+            }
+            other => Ok(other),
+        }
+    }
+}
+
 macro_rules! impl_query_output {
     ( $($ident:ty),+ $(,)? ) => { $(
         impl QueryOutput for $ident {
diff --git a/core/test_network/src/lib.rs b/core/test_network/src/lib.rs
index 7609d765742..93d8e724df4 100644
--- a/core/test_network/src/lib.rs
+++ b/core/test_network/src/lib.rs
@@ -892,15 +892,17 @@ impl TestClient for Client {
         <R::Output as QueryOutput>::Target: core::fmt::Debug,
         <R::Output as TryFrom<QueryOutputBox>>::Error: Into<eyre::Error>,
     {
-        let mut query_result = None;
         for _ in 0..max_attempts {
-            query_result = match self.request(request.clone()) {
-                Ok(result) if f(result.clone()) => return Ok(()),
-                result => Some(result),
+            if let Ok(result) = self.request(request.clone()) {
+                if f(result) {
+                    return Ok(());
+                }
             };
             thread::sleep(period);
         }
-        Err(eyre::eyre!("Failed to wait for query request completion that would satisfy specified closure. Got this query result instead: {:?}", &query_result))
+        Err(eyre::eyre!(
+            "Failed to wait for query request completion that would satisfy specified closure"
+        ))
     }
 
     fn poll_request<R: Query + Debug + Clone>(