diff --git a/src/pg/test/sql/06_segmentation_test.sql b/src/pg/test/sql/06_segmentation_test.sql
index fb6bd2b..2e93afa 100644
--- a/src/pg/test/sql/06_segmentation_test.sql
+++ b/src/pg/test/sql/06_segmentation_test.sql
@@ -1,5 +1,5 @@
 \pset format unaligned
-\set ECHO all
+\set ECHO none
 \i test/fixtures/ml_values.sql
 SELECT cdb_crankshaft._cdb_random_seeds(1234);
-SELECT prediction from cdb_crankshaft.CDB_CreateAndPredictSegment('select target,x1,x2,x3 from ml_values where class= $$train$$','target','select cartodb_id, target,x1,x2,x3 from ml_values where class=$$test$$') limit 20
+SELECT prediction FROM cdb_crankshaft.CDB_CreateAndPredictSegment('SELECT target, x1, x2, x3 FROM ml_values WHERE class = $$train$$','target','SELECT cartodb_id, target, x1, x2, x3 FROM ml_values WHERE class = $$test$$') LIMIT 20;
diff --git a/src/py/crankshaft/test/test_segmentation.py b/src/py/crankshaft/test/test_segmentation.py
index 2fc0a2d..7e806ff 100644
--- a/src/py/crankshaft/test/test_segmentation.py
+++ b/src/py/crankshaft/test/test_segmentation.py
@@ -23,6 +23,11 @@ class SegmentationTest(unittest.TestCase):
         else:
             return  [dict( zip(['x1','x2','x3','target', 'cartodb_id'],[x1,x2,x3,y,cartodb_id]))]
 
+    def test_replace_nan_with_mean(self):
+        test_array = np.array([1.2, np.nan, 3.2, np.nan, np.nan])
+
+
+
     def test_create_and_predict_segment(self):
         n_samples = 1000