def _gremlin_batch_test(_mock_gremlin, size: int):
packages = _get_normalized_packages()
_mock_gremlin.return_value = None
with mock.patch('src.v2.stack_aggregator.GREMLIN_QUERY_SIZE', size):
_mock_gremlin.reset_mock()
sa.Aggregator(
request=StackAggregatorRequest(
registration_status="REGISTERED",
uuid="3fa85f64-5717-4562-b3fc-2c963f66afa6",
external_request_id='test_request_id',
ecosystem='pypi',
manifest_file_path='/tmp/bin',
packages=[_DJANGO]),
normalized_packages=packages).get_package_details_from_graph()
ith = 0
last = 0
for i, call in enumerate(_mock_gremlin.call_args_list, start=1):
args, kwargs = call
assert len(args) == 2
assert args[0].startswith('epv')
assert isinstance(args[1], dict)
if i == len(_mock_gremlin.call_args_list):
last = len(args[1]['packages'])
else:
ith = len(args[1]['packages'])
return _mock_gremlin.call_count, ith, last
def test_gremlin_batch_call(_mock_gremlin):
"""Test post_gremlin call according to batch size."""
# empty
_mock_gremlin.return_value = None
packages = NormalizedPackages([], 'pypi')
result = sa.Aggregator(request=StackAggregatorRequest(registration_status="REGISTERED",
uuid="3fa85f64-5717-4562-b3fc-2c963f66afa6",
external_request_id='test_request_id',
ecosystem='pypi', manifest_file_path='/tmp/bin', packages=[_DJANGO]),
normalized_packages=packages).get_package_details_from_graph()
assert result is not None
assert isinstance(result, dict)
_mock_gremlin.assert_not_called()
assert (1, 0, 5) == _gremlin_batch_test(_mock_gremlin, 100)
assert (5, 1, 1) == _gremlin_batch_test(_mock_gremlin, 1)
assert (3, 2, 1) == _gremlin_batch_test(_mock_gremlin, 2)
assert (2, 3, 2) == _gremlin_batch_test(_mock_gremlin, 3)
assert (2, 4, 1) == _gremlin_batch_test(_mock_gremlin, 4)
assert (1, 0, 5) == _gremlin_batch_test(_mock_gremlin, 5)