#!/usr/bin/env python

import csv
import os
import shutil
import tempfile
import unittest
from collections import OrderedDict
from orion.columns import read_columns
from orion.transpose import transpose
from StringIO import StringIO


class TestColumns(unittest.TestCase):

    # fake data
    header = ["Animal", "Baby", "Group"]
    data = [("Fox", "Kit", "Leash"),
            ("Whale", "Calf", "Pod"),
            ("Cuttlefish", "Hatchling", "Bob"),
            ("Crow", "Hatchling", "Murder")]

    def create_data(self, fp):
        """write simple data to a test file and return its path"""

        writer = csv.writer(fp)
        writer.writerow(self.header)
        writer.writerows(self.data)
        fp.flush()

    def fake_data_assertions(self, columns):
        """test the columns created from our fake data"""

        assert len(columns) == len(self.header)
        assert set(columns.keys()) == set(self.header)
        transposed = transpose(self.data)
        assert set([len(row) for row in transposed]) == set([len(self.data)])
        for key, column in zip(self.header, transposed):
            assert len(columns[key]) == len(self.data)
            assert columns[key] == column

    def test_string(self):
        """basic test of columns interface"""

        # write fake data
        fileobj = StringIO()
        self.create_data(fileobj)
        fileobj.seek(0)

        # read into columns
        columns = read_columns(fileobj)

        # test what we have
        self.fake_data_assertions(columns)

    def test_file(self):
        """write to a real file and use that"""

        tmpdir = tempfile.mkdtemp()
        try:
            # write the data
            path = os.path.join(tmpdir, 'tmp.csv')
            with open(path, 'w') as f:
                self.create_data(f)
            assert os.path.exists(path)

            # read into columns
            columns = read_columns(path)

            # test what we have
            self.fake_data_assertions(columns)
        finally:
            shutil.rmtree(tmpdir, ignore_errors=True)

    def test_ordered_dict(self):
        """ensure we can make use of an OrderedDict"""

        # write fake data
        fileobj = StringIO()
        self.create_data(fileobj)
        fileobj.seek(0)

        # read into columns
        columns = read_columns(fileobj, type=OrderedDict)

        # test what we have
        self.fake_data_assertions(columns)


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