diff --git a/src/pyshark/capture/capture.py b/src/pyshark/capture/capture.py index 31fcd2a4..cb505504 100644 --- a/src/pyshark/capture/capture.py +++ b/src/pyshark/capture/capture.py @@ -1,4 +1,5 @@ import asyncio +import inspect import os import threading import subprocess @@ -291,7 +292,10 @@ async def _go_through_packets_from_fd(self, fd, packet_callback, packet_count=No if packet: packets_captured += 1 try: - packet_callback(packet) + if inspect.iscoroutinefunction(packet_callback): + await packet_callback(packet) + else: + packet_callback(packet) except StopCapture: self._log.debug("User-initiated capture stop in callback") break diff --git a/src/setup.py b/src/setup.py index bd027a68..f98a57e5 100644 --- a/src/setup.py +++ b/src/setup.py @@ -6,7 +6,7 @@ setup( name="pyshark", - version="0.6", + version="0.6.1", packages=find_packages(), package_data={'': ['*.ini', '*.pcapng']}, install_requires=['lxml', 'termcolor', 'packaging', 'appdirs'], diff --git a/tests/test_cap_operations.py b/tests/test_cap_operations.py index e80904e0..a6670f1f 100644 --- a/tests/test_cap_operations.py +++ b/tests/test_cap_operations.py @@ -17,6 +17,14 @@ def test_packet_callback_called_for_each_packet(lazy_simple_capture): assert mock_callback.call_count == 24 +def test_async_packet_callback_called_for_each_packet(lazy_simple_capture): + # Test cap has 24 packets + mock_callback = mock.AsyncMock() + lazy_simple_capture.apply_on_packets(mock_callback) + assert mock_callback.call_count == 24 + mock_callback.assert_awaited() + + def test_apply_on_packet_stops_on_timeout(lazy_simple_capture): def wait(pkt): time.sleep(5)