import unittest
import sys
sys.path.append('../../system')
from dataset import *
from scipy import rand,shape

class DatasetTest(unittest.TestCase):
    def setUp(self):
        self.dim_0_ids = ['sample_a','sample_b']
        self.dim_1_ids = ['gene_a','gene_b','gene_c']
        self.dim_labels = ['samples','genes']
        self.def_list = [[self.dim_labels[0],self.dim_0_ids],[self.dim_labels[1],self.dim_1_ids]]
        self.array = rand(2,3)
        self.testdata = Dataset(self.array,self.def_list)

    def testCreation(self):
        assert self.testdata._data == self.array
        assert 'sample_a' in self.testdata.ids['samples'].keys()
        assert 'gene_b' in self.testdata.ids['genes'].keys()

        
        
    def testExtraction(self):
        ids = ['gene_a','gene_b']
        dim_name = 'genes'
        subset = self.testdata.extract_data(ids,dim_name)
        assert shape(subset._data) == (2,2)
        assert subset.ids[dim_name].keys() == ids
        assert subset.ids[dim_name].values() == [0,1]


if __name__ == '__main__':
    unittest.main()