From 37d90543f4a485c46f36b0e9da42f160696974c6 Mon Sep 17 00:00:00 2001
From: clemahieu <clemahieu@gmail.com>
Date: Fri, 10 Feb 2017 19:15:17 -0600
Subject: [PATCH] Enabling work_peer tests.

---
 rai/core_test/rpc.cpp | 42 +++++++++++++++++++++++++++++++++---------
 rai/node/node.cpp     | 24 +++++++++++++++++-------
 rai/node/node.hpp     |  1 +
 3 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/rai/core_test/rpc.cpp b/rai/core_test/rpc.cpp
index da279762..d7a15271 100644
--- a/rai/core_test/rpc.cpp
+++ b/rai/core_test/rpc.cpp
@@ -1375,7 +1375,7 @@ TEST (rpc, work_cancel)
 	thread.join ();
 }
 
-TEST (rpc, DISABLED_work_peer_bad)
+TEST (rpc, work_peer_bad)
 {
     rai::system system (24000, 2);
 	rai::node_init init1;
@@ -1388,11 +1388,18 @@ TEST (rpc, DISABLED_work_peer_bad)
 	rpc.start ();
 	node2.config.work_peers.push_back (std::make_pair (boost::asio::ip::address_v6::any (), 0));
 	rai::block_hash hash1 (1);
-	auto work (node2.generate_work (hash1));
-	ASSERT_FALSE (system.work.work_validate (hash1, work));
+	uint64_t work;
+	node2.generate_work (hash1, [&work] (uint64_t work_a)
+	{
+		work = work_a;
+	});
+	while (system.work.work_validate (hash1, work))
+	{
+		system.poll ();
+	}
 }
 
-TEST (rpc, DISABLED_work_peer_one)
+TEST (rpc, work_peer_one)
 {
     rai::system system (24000, 2);
 	rai::node_init init1;
@@ -1405,11 +1412,18 @@ TEST (rpc, DISABLED_work_peer_one)
 	rpc.start ();
 	node2.config.work_peers.push_back (std::make_pair (node1.network.endpoint ().address (), rpc.config.port));
 	rai::keypair key1;
-	auto work (node2.generate_work (key1.pub));
-	ASSERT_FALSE (system.work.work_validate (key1.pub, work));
+	uint64_t work;
+	node2.generate_work (key1.pub, [&work] (uint64_t work_a)
+	{
+		work = work_a;
+	});
+	while (system.work.work_validate (key1.pub, work))
+	{
+		system.poll ();
+	}
 }
 
-TEST (rpc, DISABLED_work_peer_many)
+TEST (rpc, work_peer_many)
 {
     rai::system system1 (24000, 1);
     rai::system system2 (24001, 1);
@@ -1439,8 +1453,18 @@ TEST (rpc, DISABLED_work_peer_many)
 	for (auto i (0); i < 10; ++i)
 	{
 		rai::keypair key1;
-		auto work (node1.generate_work (key1.pub));
-		ASSERT_FALSE (system1.work.work_validate (key1.pub, work));
+		uint64_t work;
+		node1.generate_work (key1.pub, [&work] (uint64_t work_a)
+		{
+			work = work_a;
+		});
+		while (system1.work.work_validate (key1.pub, work))
+		{
+			system1.poll ();
+			system2.poll ();
+			system3.poll ();
+			system4.poll ();
+		}
 	}
 }
 
diff --git a/rai/node/node.cpp b/rai/node/node.cpp
index 501855b7..134ffd9c 100755
--- a/rai/node/node.cpp
+++ b/rai/node/node.cpp
@@ -1448,7 +1448,8 @@ boost::asio::ip::tcp::socket socket;
 class distributed_work : public std::enable_shared_from_this <distributed_work>
 {
 public:
-distributed_work (std::shared_ptr <rai::node> const & node_a, rai::block_hash const & root_a) :
+distributed_work (std::shared_ptr <rai::node> const & node_a, rai::block_hash const & root_a, std::function <void (uint64_t)> callback_a) :
+callback (callback_a),
 node (node_a),
 root (root_a)
 {
@@ -1608,7 +1609,7 @@ void set_once (uint64_t work_a)
 {
 	if (!completed.test_and_set ())
 	{
-		promise.set_value (work_a);
+		callback (work_a);
 	}
 }
 void failure (boost::asio::ip::address const & address)
@@ -1622,7 +1623,7 @@ void handle_failure (bool last)
 	{
 		if (!completed.test_and_set ())
 		{
-			promise.set_value (node->work.generate (root));
+			callback (node->work.generate (root));
 		}
 	}
 }
@@ -1632,7 +1633,7 @@ bool remove (boost::asio::ip::address const & address)
 	outstanding.erase (address);
 	return outstanding.empty ();
 }
-std::promise <uint64_t> promise;
+std::function <void (uint64_t)> callback;
 std::shared_ptr <rai::node> node;
 rai::block_hash root;
 std::mutex mutex;
@@ -1646,11 +1647,20 @@ void rai::node::generate_work (rai::block & block_a)
     block_a.block_work_set (generate_work (block_a.root ()));
 }
 
+void rai::node::generate_work (rai::uint256_union const & hash_a, std::function <void (uint64_t)> callback_a)
+{
+	auto work_generation (std::make_shared <distributed_work> (shared (), hash_a, callback_a));
+	work_generation->start ();
+}
+
 uint64_t rai::node::generate_work (rai::uint256_union const & hash_a)
 {
-	auto work_generation (std::make_shared <distributed_work> (shared (), hash_a));
-	work_generation->start ();
-	return work_generation->promise.get_future ().get ();
+	std::promise <uint64_t> promise;
+	generate_work (hash_a, [&promise] (uint64_t work_a)
+	{
+		promise.set_value (work_a);
+	});
+	return promise.get_future ().get ();
 }
 
 namespace
diff --git a/rai/node/node.hpp b/rai/node/node.hpp
index 4cbf140c..bd07834c 100644
--- a/rai/node/node.hpp
+++ b/rai/node/node.hpp
@@ -362,6 +362,7 @@ public:
 	int price (rai::uint128_t const &, int);
 	void generate_work (rai::block &);
 	uint64_t generate_work (rai::uint256_union const &);
+	void generate_work (rai::uint256_union const &, std::function <void (uint64_t)>);
 	bool rollback_predicate (rai::block const &);
 	rai::node_config config;
     rai::alarm & alarm;