4 files changed, 25 insertions, 6 deletions
diff --git a/tests/block_weight/CMakeLists.txt b/tests/block_weight/CMakeLists.txt
index b0d716ea0..a59c9c1fc 100644
--- a/tests/block_weight/CMakeLists.txt
+++ b/tests/block_weight/CMakeLists.txt
@@ -42,4 +42,4 @@ target_link_libraries(block_weight
 
 add_test(
   NAME    block_weight
-  COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/compare.py ${CMAKE_CURRENT_SOURCE_DIR}/block_weight.py ${CMAKE_CURRENT_BINARY_DIR}/block_weight)
+  COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/compare.py ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/block_weight.py ${CMAKE_CURRENT_BINARY_DIR}/block_weight)
diff --git a/tests/block_weight/block_weight.cpp b/tests/block_weight/block_weight.cpp
index 57fcb497e..1e8974b78 100644
--- a/tests/block_weight/block_weight.cpp
+++ b/tests/block_weight/block_weight.cpp
@@ -72,6 +72,18 @@ public:
   virtual uint64_t height() const override { return blocks.size(); }
   virtual size_t get_block_weight(const uint64_t &h) const override { return blocks[h].weight; }
   virtual uint64_t get_block_long_term_weight(const uint64_t &h) const override { return blocks[h].long_term_weight; }
+  virtual std::vector<uint64_t> get_block_weights(uint64_t start_height, size_t count) const override {
+    std::vector<uint64_t> ret;
+    ret.reserve(count);
+    while (count-- && start_height < blocks.size()) ret.push_back(blocks[start_height++].weight);
+    return ret;
+  }
+  virtual std::vector<uint64_t> get_long_term_block_weights(uint64_t start_height, size_t count) const override {
+    std::vector<uint64_t> ret;
+    ret.reserve(count);
+    while (count-- && start_height < blocks.size()) ret.push_back(blocks[start_height++].long_term_weight);
+    return ret;
+  }
   virtual crypto::hash top_block_hash(uint64_t *block_height = NULL) const override {
     uint64_t h = height();
     crypto::hash top = crypto::null_hash;
diff --git a/tests/block_weight/block_weight.py b/tests/block_weight/block_weight.py
index 06aaabb02..d6fd494e3 100755
--- a/tests/block_weight/block_weight.py
+++ b/tests/block_weight/block_weight.py
@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 # Simulate a maximal block attack on the Monero network
 # This uses the scheme proposed by ArticMine
 # Written by Sarang Nother
diff --git a/tests/block_weight/compare.py b/tests/block_weight/compare.py
index c6be05206..e058a7079 100755
--- a/tests/block_weight/compare.py
+++ b/tests/block_weight/compare.py
@@ -3,10 +3,17 @@
 import sys
 import subprocess
 
-print 'running: ', sys.argv[1]
-S0 = subprocess.check_output(sys.argv[1], stderr=subprocess.STDOUT)
-print 'running: ', sys.argv[2]
-S1 = subprocess.check_output(sys.argv[2], stderr=subprocess.STDOUT)
+if len(sys.argv) == 4:
+  first = [sys.argv[1], sys.argv[2]]
+  second = [sys.argv[3]]
+else:
+  first = [sys.argv[1]]
+  second = [sys.argv[2]]
+
+print 'running: ', first
+S0 = subprocess.check_output(first, stderr=subprocess.STDOUT)
+print 'running: ', second
+S1 = subprocess.check_output(second, stderr=subprocess.STDOUT)
 print 'comparing'
 if S0 != S1:
   sys.exit(1)