"""
Tests for molnet function 
"""
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

import os
import unittest
import numpy as np
import pandas as pd
import deepchem as dc
import csv


class TestMolnet(unittest.TestCase):
  """
  Test basic function of molnet
  """

  def setUp(self):
    super(TestMolnet, self).setUp()
    self.current_dir = os.path.dirname(os.path.abspath(__file__))

  def test_delaney_graphconvreg(self):
    """Tests molnet benchmarking on delaney with graphconvreg."""
    datasets = ['delaney']
    model = 'graphconvreg'
    split = 'random'
    out_path = self.current_dir
    dc.molnet.run_benchmark(datasets, str(model), split=split, out_path=out_path)
    with open(os.path.join(out_path, 'results.csv'), 'r') as f:
      reader = csv.reader(f)
      for lastrow in reader:
        pass
      assert lastrow[-4] == model
      assert lastrow[-5] == 'valid'
      assert float(lastrow[-3]) > 0.75
    os.remove(os.path.join(out_path, 'results.csv'))

  def test_qm7_multitask(self):
    """Tests molnet benchmarking on qm7 with multitask network."""
    datasets = ['qm7']
    model = 'tf_regression'
    split = 'random'
    out_path = self.current_dir
    dc.molnet.run_benchmark(datasets, str(model), split=split, out_path=out_path)
    with open(os.path.join(out_path, 'results.csv'), 'r') as f:
      reader = csv.reader(f)
      for lastrow in reader:
        pass
      assert lastrow[-4] == model + '_ft'
      assert lastrow[-5] == 'valid'
      assert float(lastrow[-3]) > 0.95
    os.remove(os.path.join(out_path, 'results.csv'))

  def test_tox21_multitask(self):
    """Tests molnet benchmarking on tox21 with multitask network."""
    datasets = ['tox21']
    model = 'tf'
    split = 'random'
    out_path = self.current_dir
    dc.molnet.run_benchmark(datasets, str(model), split=split, out_path=out_path)
    with open(os.path.join(out_path, 'results.csv'), 'r') as f:
      reader = csv.reader(f)
      for lastrow in reader:
        pass
      assert lastrow[-4] == model
      assert lastrow[-5] == 'valid'
      assert float(lastrow[-3]) > 0.75
    os.remove(os.path.join(out_path, 'results.csv'))


@@ -675,32 +675,41 @@
 t.run_benchmark(
+%0A        
 datasets, str(mo
@@ -738,32 +738,32 @@
 _path=out_path)%0A
-
     with open(os
@@ -1267,32 +1267,41 @@
 t.run_benchmark(
+%0A        
 datasets, str(mo
@@ -1870,16 +1870,25 @@
 nchmark(
+%0A        
 datasets
@@ -1923,32 +1923,43 @@
 ut_path=out_path
+, test=True
 )%0A    with open(
@@ -2134,37 +2134,36 @@
 lastrow%5B-5%5D == '
-valid
+test
 '%0A      assert f
@@ -2177,33 +2177,32 @@
 strow%5B-3%5D) %3E 0.7
-5
 %0A    os.remove(o
@@ -2219,29 +2219,28 @@
 n(out_path, 'results.csv'))%0A
-%0A

fix style