initial dataset batch prediction

2 files changed, 14 insertions, 2 deletions

diff --git a/test/model-classification.rb b/test/model-classification.rb
index 7751bba..ca6eb27 100644
--- a/test/model-classification.rb
+++ b/test/model-classification.rb
@@ -94,6 +94,18 @@ class LazarClassificationTest < MiniTest::Test
     training_dataset.delete
   end
 
+  def test_dataset_prediction
+    training_dataset = Dataset.from_csv_file File.join(DATA_DIR,"hamster_carcinogenicity.csv")
+    model = Model::Lazar.create training_dataset: training_dataset
+    result = model.predict training_dataset
+    assert 3, result.features.size
+    assert 8, result.compounds.size
+    assert_equal ["true"], result.values(result.compounds.first, result.features[0])
+    assert_equal [0.65], result.values(result.compounds.first, result.features[1])
+    assert_equal [0], result.values(result.compounds.first, result.features[2]) # classification returns nil, check if 
+    #p prediction_dataset
+  end
+
   def test_carcinogenicity_rf_classification
     skip "Caret rf may run into a (endless?) loop for some compounds."
     dataset = Dataset.from_csv_file "#{DATA_DIR}/multi_cell_call.csv"
diff --git a/test/setup.rb b/test/setup.rb
index c4c04cb..51871a2 100644
--- a/test/setup.rb
+++ b/test/setup.rb
@@ -3,8 +3,8 @@ require 'minitest/autorun'
 require_relative '../lib/lazar.rb'
 #require 'lazar'
 include OpenTox
-#$mongo.database.drop
-#$gridfs = $mongo.database.fs # recreate GridFS indexes
+$mongo.database.drop
+$gridfs = $mongo.database.fs # recreate GridFS indexes
 #PhysChem.descriptors
 TEST_DIR ||= File.expand_path(File.dirname(__FILE__))
 DATA_DIR ||= File.join(TEST_DIR,"data")