def test_collection_items(mocker, count, start, filters, expected_count):
mocker.patch(
"arche.tools.api.get_collection",
return_value=Collection(len(collection_items)),
autospec=True,
)
get_items_mock = mocker.patch(
"arche.tools.api.get_items",
return_value=collection_items[:expected_count],
autospec=True,
)
items = CollectionItems("key", count, start, filters)
assert items.key == "key"
assert items.start == start
assert items.filters == filters
np.testing.assert_array_equal(items.raw, collection_items[:expected_count])
pd.testing.assert_frame_equal(items.df,
expected_col_df.iloc[:expected_count])
assert len(items) == expected_count
assert items.limit == len(collection_items)
assert items.count == expected_count
get_items_mock.assert_called_once_with("key",
expected_count,
0,
start,
filters,
desc="Fetching from 'key'")
def get_items(
source: Union[str, pd.DataFrame, RawItems],
start: int,
count: Optional[int],
filters: Optional[api.Filters],
expand: bool,
) -> Union[JobItems, CollectionItems]:
if isinstance(source, pd.DataFrame):
return Items.from_df(source, expand=expand)
elif isinstance(source, Iterable) and not isinstance(source, str):
return Items.from_array(source, expand=expand)
elif helpers.is_job_key(source):
return JobItems(key=source,
start=start,
count=count,
filters=filters,
expand=expand)
elif helpers.is_collection_key(source):
if start:
raise ValueError(
"Collections API does not support 'start' parameter")
return CollectionItems(key=source,
count=count,
filters=filters,
expand=expand)
else:
raise ValueError(
f"'{source}' is not a valid job or collection key")
def get_collection_items(request, mocker):
mocker.patch("arche.tools.api.get_collection",
return_value=get_collection,
autospec=True)
mocker.patch(
"arche.readers.items.CollectionItems.fetch_data",
return_value=np.array(request.param),
autospec=True,
)
collection_items = CollectionItems(key="112358/collections/s/pages",
count=len(request.param))
return collection_items
def get_collection_items_mock(mocker, items=default_items, **kwargs):
mocker.patch(
"arche.readers.items.CollectionItems.fetch_data",
return_value=items,
autospec=True,
)
collection_items = CollectionItems(key=kwargs.get(
"key", "112358/collections/s/pages"),
count=kwargs.get("count", len(items)),
filters=kwargs.get("filters", None),
expand=kwargs.get("expand", True),
**kwargs)
return collection_items
def get_items(
source: Union[str, pd.DataFrame, RawItems],
count: Optional[int],
start: Optional[str],
filters: Optional[api.Filters],
) -> Items:
if isinstance(source, pd.DataFrame):
return Items.from_df(source)
elif isinstance(source, Iterable) and not isinstance(source, str):
return Items.from_array(cast(RawItems, source))
elif helpers.is_job_key(source):
return JobItems(source, count, int(start or 0), filters)
elif helpers.is_collection_key(source):
return CollectionItems(source, count, start, filters)
else:
raise ValueError(
f"'{source}' is not a valid job or collection key")
def get_items(
source: str,
start: int,
count: Optional[int],
filters: Optional[api.Filters],
expand: bool,
) -> Union[JobItems, CollectionItems]:
if helpers.is_job_key(source):
return JobItems(
key=source, start=start, count=count, filters=filters, expand=expand
)
elif helpers.is_collection_key(source):
if start:
raise ValueError("Collections API does not support 'start' parameter")
return CollectionItems(
key=source, count=count, filters=filters, expand=expand
)
else:
raise ValueError(f"'{source}' is not a valid job or collection key")
def test_collection_items(mocker, count, filters, expand, expected_count):
mocker.patch(
"arche.tools.api.get_collection",
return_value=Collection(collection_items),
autospec=True,
)
get_items_mock = mocker.patch(
"arche.tools.api.get_items",
return_value=collection_items[:expected_count],
autospec=True,
)
items = CollectionItems("key", count, filters, expand)
assert items.key == "key"
assert items._count == count
assert items.filters == filters
assert items.expand == expand
assert items.dicts == collection_items[:expected_count]
assert items._flat_df is None
assert items.limit == len(collection_items)
assert items.count == expected_count
get_items_mock.assert_called_once_with("key", expected_count, 0, filters)
def test_collection_items(mocker, count, filters, expand, expected_count):
mocker.patch(
"arche.tools.api.get_collection",
return_value=Collection(len(collection_items)),
autospec=True,
)
get_items_mock = mocker.patch(
"arche.tools.api.get_items",
return_value=collection_items[:expected_count],
autospec=True,
)
items = CollectionItems("key", count, filters, expand)
assert items.key == "key"
assert items.filters == filters
assert items.expand == expand
np.testing.assert_array_equal(items.raw, collection_items[:expected_count])
pd.testing.assert_frame_equal(items.df,
expected_col_items.iloc[:expected_count])
pd.testing.assert_frame_equal(items.flat_df, items.df)
assert len(items) == len(expected_col_items.iloc[:expected_count])
assert items.limit == len(collection_items)
assert items.count == expected_count
get_items_mock.assert_called_once_with("key", expected_count, 0, filters)