1 files changed, 9 insertions, 4 deletions
diff --git a/klm/lm/builder/adjust_counts_test.cc b/klm/lm/builder/adjust_counts_test.cc
index 68b5f33e..073c5dfe 100644
--- a/klm/lm/builder/adjust_counts_test.cc
+++ b/klm/lm/builder/adjust_counts_test.cc
@@ -1,6 +1,6 @@
 #include "lm/builder/adjust_counts.hh"
 
-#include "lm/builder/multi_stream.hh"
+#include "lm/builder/ngram_stream.hh"
 #include "util/scoped.hh"
 
 #include <boost/thread/thread.hpp>
@@ -61,19 +61,24 @@ BOOST_AUTO_TEST_CASE(Simple) {
     util::stream::ChainConfig config;
     config.total_memory = 100;
     config.block_count = 1;
-    Chains chains(4);
+    util::stream::Chains chains(4);
     for (unsigned i = 0; i < 4; ++i) {
       config.entry_size = NGram::TotalSize(i + 1);
       chains.push_back(config);
     }
 
     chains[3] >> WriteInput();
-    ChainPositions for_adjust(chains);
+    util::stream::ChainPositions for_adjust(chains);
     for (unsigned i = 0; i < 4; ++i) {
       chains[i] >> boost::ref(outputs[i]);
     }
     chains >> util::stream::kRecycle;
-    BOOST_CHECK_THROW(AdjustCounts(counts, discount).Run(for_adjust), BadDiscountException);
+    std::vector<uint64_t> counts_pruned(4);
+    std::vector<uint64_t> prune_thresholds(4);
+    DiscountConfig discount_config;
+    discount_config.fallback = Discount();
+    discount_config.bad_action = THROW_UP;
+    BOOST_CHECK_THROW(AdjustCounts(prune_thresholds, counts, counts_pruned, discount_config, discount).Run(for_adjust), BadDiscountException);
   }
   BOOST_REQUIRE_EQUAL(4UL, counts.size());
   BOOST_CHECK_EQUAL(4UL, counts[0]);