def test__page(self):
mock_api = MagicMock()
mock_api.list_dataset_items.return_value = self._build_dataset_items_response()
iterator = DatasetItemIterator(
mock_api, self.organization_id, self.dataset_id)
page = iterator._page()
self.assertIsInstance(page, list)
item = page[0]
self.assertIsInstance(item, DatasetItem)
def test__page_iter(self):
mock_api = MagicMock()
mock_api.list_dataset_items.side_effect = [
self._build_dataset_items_response(),
self._build_empty_dataset_items_response()
]
iterator = DatasetItemIterator(
mock_api, self.organization_id, self.dataset_id)
page_iterator = iterator._page_iter()
for page in page_iterator:
self.assertIsInstance(page, list)
for item in page:
self.assertIsInstance(item, DatasetItem)
def test_next_raise_stop_iteration(self):
mock_api = MagicMock()
mock_api.list_dataset_items.side_effect = [
{
'next_page_token': 'dummy1',
'items': [
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_1'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_2'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_3'},
]
},
{
'next_page_token': 'dummy2',
'items': [
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_4'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_5'}
]
},
{
'next_page_token': None,
'items': []
}
]
iterator = DatasetItemIterator(
mock_api,
organization_id=self.organization_id,
dataset_id=self.dataset_id)
for i in range(5):
next(iterator)
with self.assertRaises(StopIteration):
next(iterator)
self.assertEqual(mock_api.list_dataset_items.call_count, 3)
def test__items_iter(self):
mock_api = MagicMock()
mock_api.list_dataset_items.side_effect = [
self._build_dataset_items_response(),
self._build_empty_dataset_items_response()
]
iterator = DatasetItemIterator(
mock_api, self.organization_id, self.dataset_id)
with patch('abeja.datalake.file.DatalakeFile.get_content') as m:
for item in iterator:
self.assertIsInstance(item, DatasetItem)
self.assertEqual(m.call_count, 0)
def test_next_up_to_next_page_and_iter_1(self):
mock_api = MagicMock()
mock_api.list_dataset_items.side_effect = [
{
'next_page_token': 'dummy1',
'items': [
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_1'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_2'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_3'},
]
},
{
'next_page_token': 'dummy2',
'items': [
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_4'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_5'}
]
},
{
'next_page_token': None,
'items': []
}
]
iterator = DatasetItemIterator(
mock_api,
organization_id=self.organization_id,
dataset_id=self.dataset_id)
item_1 = next(iterator)
self.assertEqual(item_1.dataset_item_id, 'item_id_1')
item_2 = next(iterator)
self.assertEqual(item_2.dataset_item_id, 'item_id_2')
item_3 = next(iterator)
self.assertEqual(item_3.dataset_item_id, 'item_id_3')
mock_api.list_dataset_items.assert_called_with(
self.organization_id, self.dataset_id, params={})
item_4 = next(iterator)
self.assertEqual(item_4.dataset_item_id, 'item_id_4')
# take last one from __iter__
self.assertEqual(len(list(iterator)), 1)
self.assertEqual(mock_api.list_dataset_items.call_count, 3)
def test_next(self):
mock_api = MagicMock()
mock_api.list_dataset_items.side_effect = [
{
'next_page_token': 'dummy1',
'items': [
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_1'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_2'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_3'},
]
},
{
'next_page_token': 'dummy2',
'items': [
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_4'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_5'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_6'},
]
},
{
'next_page_token': 'dummy3',
'items': [
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_7'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_8'},
{'dataset_id': self.dataset_id, 'dataset_item_id': 'item_id_9'},
]
},
{
'next_page_token': None,
'items': []
}
]
iterator = DatasetItemIterator(
mock_api,
organization_id=self.organization_id,
dataset_id=self.dataset_id)
item_1 = next(iterator)
self.assertEqual(item_1.dataset_item_id, 'item_id_1')
item_2 = next(iterator)
self.assertEqual(item_2.dataset_item_id, 'item_id_2')
self.assertEqual(len(list(iterator)), 7)