class TestFrame(TestBase): def setUp(self): TestBase.setUp(self) self.dframe = self.get_data('good_eats.csv') self.bframe = BambooFrame(self.dframe) def _add_bamboo_reserved_keys(self, value=1): for key in BAMBOO_RESERVED_KEYS: column = Series([value] * len(self.bframe)) column.name = key self.bframe = BambooFrame(self.bframe.join(column)) def test_add_parent_column(self): value = 1 self._add_bamboo_reserved_keys(value) for index, item in self.bframe[PARENT_DATASET_ID].iteritems(): self.assertEqual(item, value) def test_decode_mongo_reserved_keys(self): prev_columns = self.bframe.columns for col in MONGO_RESERVED_KEYS: self.assertTrue(col in self.bframe.columns) self.bframe.decode_mongo_reserved_keys() for col in MONGO_RESERVED_KEYS: self.assertFalse(col in self.bframe.columns) def test_recognize_dates(self): bframe_with_dates = self.bframe.recognize_dates() for field in bframe_with_dates['submit_date']: self.assertTrue(isinstance(field, datetime)) def test_recognize_dates_from_schema(self): schema = Schema({ 'submit_date': { SIMPLETYPE: DATETIME } }) bframe_with_dates = self.bframe.recognize_dates_from_schema(schema) for field in bframe_with_dates['submit_date']: self.assertTrue(isinstance(field, datetime)) def test_remove_bamboo_reserved_keys(self): self._add_bamboo_reserved_keys() for key in BAMBOO_RESERVED_KEYS: self.assertTrue(key in self.bframe.columns) self.bframe.remove_bamboo_reserved_keys() for key in BAMBOO_RESERVED_KEYS: self.assertFalse(key in self.bframe.columns) def test_remove_bamboo_reserved_keys_exclusion(self): self._add_bamboo_reserved_keys() for key in BAMBOO_RESERVED_KEYS: self.assertTrue(key in self.bframe.columns) self.bframe.remove_bamboo_reserved_keys([PARENT_DATASET_ID]) for key in BAMBOO_RESERVED_KEYS: if key == PARENT_DATASET_ID: self.assertTrue(key in self.bframe.columns) else: self.assertFalse(key in self.bframe.columns) def test_only_rows_for_parent_id(self): parent_id = 1 len_parent_rows = len(self.bframe) / 2 column = Series([parent_id] * len_parent_rows) column.name = PARENT_DATASET_ID self.bframe = BambooFrame(self.bframe.join(column)) bframe_only = self.bframe.only_rows_for_parent_id(parent_id) self.assertFalse(PARENT_DATASET_ID in bframe_only.columns) self.assertEqual(len(bframe_only), len_parent_rows) def test_to_jsondict(self): jsondict = self.bframe.to_jsondict() self.assertEqual(len(jsondict), len(self.bframe)) for col in jsondict: self.assertEqual(len(col), len(self.bframe.columns)) def test_to_json(self): json = self.bframe.to_json() self.assertEqual(type(json), str)