1 files changed, 9 insertions, 3 deletions

diff --git a/kad/tests/searchCancel.hs b/kad/tests/searchCancel.hs
index 85986674..33860a2f 100644
--- a/kad/tests/searchCancel.hs
+++ b/kad/tests/searchCancel.hs
@@ -10,11 +10,13 @@ import Network.Kademlia.Persistence
 import Network.Kademlia.Routing
 import Network.Kademlia.Search
 
+import Network.QueryResponse as QR
 import qualified Data.MinMaxPSQ as MM
 import qualified Data.Set as Set
 
-makeSchResults :: TVar (Maybe (SearchState Int Int () Int Int)) -> TVar Int -> IO (Maybe ([Int],[Int],Maybe ()))
+makeSchResults :: TVar (Maybe (SearchState Int Int () Int Int Int)) -> TVar Int -> IO (Maybe ([Int],[Int],Maybe ()))
 makeSchResults mbv var = do
+    putStrLn "makeSchResults"
     threadDelay 200000
     (r,io) <- atomically $ do
         n <- readTVar var
@@ -42,11 +44,15 @@ kad = KademliaSpace
     , kademliaSample   = \_ x _ -> pure x
     }
 
-sch :: TVar (Maybe (SearchState Int Int () Int Int)) -> TVar Int -> Search Int Int () Int Int
+sch :: TVar (Maybe (SearchState Int Int () Int Int Int)) -> TVar Int -> Search Int Int () Int Int Int
 sch mbv var = Search
     { searchSpace       = kad
     , searchNodeAddress = id
-    , searchQuery       = \_ _ -> makeSchResults mbv var
+    , searchQuery       = \_ _ f -> do r <- makeSchResults mbv var
+                                       let qk = maybe 0 (\(ns,_,_) -> head ns) r
+                                       f qk $ maybe TimedOut Success r
+                                       return qk
+    , searchQueryCancel = \_ -> return ()
     , searchAlpha       = 4
     , searchK           = 8
     }