# Path: bamboo/core/frame.py
# BAMBOO_RESERVED_KEYS = [
#     DATASET_ID,
#     INDEX,
#     PARENT_DATASET_ID,
# ]
# 
# def remove_reserved_keys(df, exclude=[]):
#     """Remove reserved internal columns in this DataFrame.
# 
#     :param keep_parent_ids: Keep parent column if True, default False.
#     """
#     reserved_keys = __column_intersect(
#         df, BAMBOO_RESERVED_KEYS).difference(set(exclude))
# 
#     return df.drop(reserved_keys, axis=1)
# 
# def rows_for_parent_id(df, parent_id):
#     """DataFrame with only rows for `parent_id`.
# 
#     :param parent_id: The ID to restrict rows to.
# 
#     :returns: A DataFrame including only rows with a parent ID equal to
#         that passed in.
#     """
#     return df[df[PARENT_DATASET_ID] == parent_id].drop(PARENT_DATASET_ID, 1)
# 
# PARENT_DATASET_ID = BAMBOO_RESERVED_KEY_PREFIX + 'parent_dataset_id'
# 
# Path: bamboo/tests/test_base.py
# class TestBase(unittest.TestCase):
# 
#     FIXTURE_PATH = 'tests/fixtures/'
#     SLEEP_DELAY = 0.2
#     TEST_DATASETS = [
#         'good_eats.csv',
#         'good_eats_large.csv',
#         'good_eats_with_calculations.csv',
#         'kenya_secondary_schools_2007.csv',
#         'soil_samples.csv',
#         'water_points.csv',
#         'unicode.csv',
#     ]
# 
#     test_dataset_ids = {}
# 
#     def setUp(self):
#         self.__drop_database()
#         self.__create_database()
#         self.__load_test_data()
# 
#     def tearDown(self):
#         self.__drop_database()
# 
#     def get_data(self, dataset_name):
#         return read_csv('%s%s' % (self._local_fixture_prefix(), dataset_name),
#                         encoding='utf-8')
# 
#     def _create_dataset_from_url(self, url):
#         dataset = Dataset.create()
#         return dataset.import_from_url(url, allow_local_file=True).dataset_id
# 
#     def _local_fixture_prefix(self, filename=''):
#         return 'file://localhost%s/tests/fixtures/%s' % (os.getcwd(), filename)
# 
#     def _fixture_path_prefix(self, filename=''):
#         return '/%s/tests/fixtures/%s' % (os.getcwd(), filename)
# 
#     def _file_mock(self, file_path, add_prefix=False):
#         if add_prefix:
#             file_path = self._fixture_path_prefix(file_path)
# 
#         file_ = open(file_path, 'r')
# 
#         return MockUploadedFile(file_)
# 
#     def _post_file(self, file_name='good_eats.csv'):
#         dataset = Dataset.create()
#         return dataset.import_from_csv(
#             self._file_mock(self._fixture_path_prefix(file_name))).dataset_id
# 
#     def _wait_for_dataset_state(self, dataset_id):
#         while True:
#             dataset = Dataset.find_one(dataset_id)
# 
#             if dataset.state != Dataset.STATE_PENDING:
#                 break
# 
#             sleep(self.SLEEP_DELAY)
# 
#         return dataset
# 
#     def __create_database(self):
#         Database.db(TEST_DATABASE_NAME)
# 
#     def __drop_database(self):
#         Database.client().drop_database(TEST_DATABASE_NAME)
# 
#     def __load_test_data(self):
#         for dataset_name in self.TEST_DATASETS:
#             self.test_dataset_ids[dataset_name] = uuid.uuid4().hex


from pandas import Series
from bamboo.core.frame import BAMBOO_RESERVED_KEYS, remove_reserved_keys,\
    rows_for_parent_id, PARENT_DATASET_ID
from bamboo.tests.test_base import TestBase



class TestFrame(TestBase):

    def setUp(self):
        TestBase.setUp(self)
        self.dframe = self.get_data('good_eats.csv')

    def _add_bamboo_reserved_keys(self, value=1):
        for key in BAMBOO_RESERVED_KEYS:
            column = Series([value] * len(self.dframe))
            column.name = key
            self.dframe = self.dframe.join(column)

    def test_add_parent_column(self):
        value = 1
        self._add_bamboo_reserved_keys(value)

        for index, item in self.dframe[PARENT_DATASET_ID].iteritems():
            self.assertEqual(item, value)

    def test_remove_reserved_keys(self):
        self._add_bamboo_reserved_keys()

        for key in BAMBOO_RESERVED_KEYS:
            self.assertTrue(key in self.dframe.columns)


# Path: bamboo/core/frame.py
# BAMBOO_RESERVED_KEYS = [
#     DATASET_ID,
#     INDEX,
#     PARENT_DATASET_ID,
# ]
# 
# def remove_reserved_keys(df, exclude=[]):
#     """Remove reserved internal columns in this DataFrame.
# 
#     :param keep_parent_ids: Keep parent column if True, default False.
#     """
#     reserved_keys = __column_intersect(
#         df, BAMBOO_RESERVED_KEYS).difference(set(exclude))
# 
#     return df.drop(reserved_keys, axis=1)
# 
# def rows_for_parent_id(df, parent_id):
#     """DataFrame with only rows for `parent_id`.
# 
#     :param parent_id: The ID to restrict rows to.
# 
#     :returns: A DataFrame including only rows with a parent ID equal to
#         that passed in.
#     """
#     return df[df[PARENT_DATASET_ID] == parent_id].drop(PARENT_DATASET_ID, 1)
# 
# PARENT_DATASET_ID = BAMBOO_RESERVED_KEY_PREFIX + 'parent_dataset_id'
# 
# Path: bamboo/tests/test_base.py
# class TestBase(unittest.TestCase):
# 
#     FIXTURE_PATH = 'tests/fixtures/'
#     SLEEP_DELAY = 0.2
#     TEST_DATASETS = [
#         'good_eats.csv',
#         'good_eats_large.csv',
#         'good_eats_with_calculations.csv',
#         'kenya_secondary_schools_2007.csv',
#         'soil_samples.csv',
#         'water_points.csv',
#         'unicode.csv',
#     ]
# 
#     test_dataset_ids = {}
# 
#     def setUp(self):
#         self.__drop_database()
#         self.__create_database()
#         self.__load_test_data()
# 
#     def tearDown(self):
#         self.__drop_database()
# 
#     def get_data(self, dataset_name):
#         return read_csv('%s%s' % (self._local_fixture_prefix(), dataset_name),
#                         encoding='utf-8')
# 
#     def _create_dataset_from_url(self, url):
#         dataset = Dataset.create()
#         return dataset.import_from_url(url, allow_local_file=True).dataset_id
# 
#     def _local_fixture_prefix(self, filename=''):
#         return 'file://localhost%s/tests/fixtures/%s' % (os.getcwd(), filename)
# 
#     def _fixture_path_prefix(self, filename=''):
#         return '/%s/tests/fixtures/%s' % (os.getcwd(), filename)
# 
#     def _file_mock(self, file_path, add_prefix=False):
#         if add_prefix:
#             file_path = self._fixture_path_prefix(file_path)
# 
#         file_ = open(file_path, 'r')
# 
#         return MockUploadedFile(file_)
# 
#     def _post_file(self, file_name='good_eats.csv'):
#         dataset = Dataset.create()
#         return dataset.import_from_csv(
#             self._file_mock(self._fixture_path_prefix(file_name))).dataset_id
# 
#     def _wait_for_dataset_state(self, dataset_id):
#         while True:
#             dataset = Dataset.find_one(dataset_id)
# 
#             if dataset.state != Dataset.STATE_PENDING:
#                 break
# 
#             sleep(self.SLEEP_DELAY)
# 
#         return dataset
# 
#     def __create_database(self):
#         Database.db(TEST_DATABASE_NAME)
# 
#     def __drop_database(self):
#         Database.client().drop_database(TEST_DATABASE_NAME)
# 
#     def __load_test_data(self):
#         for dataset_name in self.TEST_DATASETS:
#             self.test_dataset_ids[dataset_name] = uuid.uuid4().hex


# Path: bamboo/tests/core/test_frame.py
from pandas import Series
from bamboo.core.frame import BAMBOO_RESERVED_KEYS, remove_reserved_keys,\
    rows_for_parent_id, PARENT_DATASET_ID
from bamboo.tests.test_base import TestBase




class TestFrame(TestBase):

    def setUp(self):
        TestBase.setUp(self)
        self.dframe = self.get_data('good_eats.csv')

    def _add_bamboo_reserved_keys(self, value=1):
        for key in BAMBOO_RESERVED_KEYS:
            column = Series([value] * len(self.dframe))
            column.name = key
            self.dframe = self.dframe.join(column)

    def test_add_parent_column(self):
        value = 1
        self._add_bamboo_reserved_keys(value)

        for index, item in self.dframe[PARENT_DATASET_ID].iteritems():
            self.assertEqual(item, value)

    def test_remove_reserved_keys(self):
        self._add_bamboo_reserved_keys()

        for key in BAMBOO_RESERVED_KEYS:
            self.assertTrue(key in self.dframe.columns)


        dframe = remove_reserved_keys(self.dframe)