diff --git a/src/chweb/consumer.py b/src/chweb/consumer.py index cb18681..7da5bdb 100644 --- a/src/chweb/consumer.py +++ b/src/chweb/consumer.py @@ -19,6 +19,7 @@ class Consumer(Service): they can be serialized in a :class:`chweb.models.Check` object, then they are saved in the database. """ + def __init__(self, config: Config, logger: logging.Logger, event_loop: asyncio.AbstractEventLoop, @@ -110,14 +111,14 @@ class Db: """ if self.conn is not None: try: - await self.conn.execute(''' - INSERT INTO statuses (domain, regex, regex_matches, + await self.conn.execute( + '''INSERT INTO statuses (domain, regex, regex_matches, request_time, response_time, status, url) - VALUES($1, $2, $3, $4, $5, $6, $7) - ''', data.domain, data.regex, data.regex_matches, - data.request_time, data.response_time, data.status, - data.url) + VALUES($1, $2, $3, $4, $5, $6, $7)''', + data.domain, data.regex, data.regex_matches, + data.request_time, data.response_time, data.status, + data.url) except asyncpg.PostgresError as exc: self.logger.error("error in query %s", exc) raise diff --git a/tests/conftest.py b/tests/conftest.py index 5ddc26e..848a0b6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import asyncio import pytest from chweb.cmd import create_config +from chweb.models import Check @pytest.fixture() def config(): @@ -47,3 +48,12 @@ def config_invalid(): ] } return create_config(config_dict) + +@pytest.fixture +def check(): + return Check( + domain="example.com", + response_time=3265, + status=200, + url="https://example.com", + ) diff --git a/tests/test_consumer.py b/tests/test_consumer.py new file mode 100644 index 0000000..ed2f3d8 --- /dev/null +++ b/tests/test_consumer.py @@ -0,0 +1,21 @@ +import asyncio + +import aiokafka +from mock import Mock, AsyncMock +import pytest + +from chweb.consumer import Consumer + + +@pytest.mark.asyncio +async def test_consumer_called(check, config, event_loop): + consumer = Consumer(config, Mock(), event_loop, Mock()) + + consumer.consumer = AsyncMock() + consumer.db = AsyncMock() + + task = event_loop.create_task(consumer()) + await asyncio.sleep(0) + consumer.db.setup.assert_called() + consumer.consumer.start.assert_called() + task.cancel() diff --git a/tests/test_producer.py b/tests/test_producer.py index e71f585..8150b99 100644 --- a/tests/test_producer.py +++ b/tests/test_producer.py @@ -1,25 +1,19 @@ import asyncio import aiokafka -from mock import Mock +from mock import Mock, AsyncMock import pytest from chweb.collector import Producer -from chweb.models import Check @pytest.mark.asyncio -async def test_producer_called(config, event_loop): +async def test_producer_called(check, config, event_loop): queue = asyncio.Queue() producer = Producer(config, Mock(), event_loop, queue) - check = Check() await queue.put(check) - async def async_patch(): - pass - Mock.__await__ = lambda x: async_patch().__await__() - - producer.producer = Mock() + producer.producer = AsyncMock() task = event_loop.create_task(producer()) await asyncio.sleep(0) @@ -32,14 +26,9 @@ async def test_producer_called(config, event_loop): async def test_producer_called_invalid(config, event_loop): queue = asyncio.Queue() producer = Producer(config, Mock(), event_loop, queue) - check = Check() await queue.put('') - async def async_patch(): - pass - Mock.__await__ = lambda x: async_patch().__await__() - - producer.producer = Mock() + producer.producer = AsyncMock() task = event_loop.create_task(producer()) await asyncio.sleep(0)