Batch insert checks in the db (#1)

tests/conftest.py

@@ -21,6 +21,7 @@ def config():
             'dbname': "chweb",
             'dbuser': "vladan",
             'dbpass': "",
+            'batch_size': 3,
         },
         'sites': [{
             'url': "https://example.com",
@@ -69,3 +70,27 @@ def check():
         status=200,
         url="https://example.com",
     )
+
+
+@pytest.fixture
+def three_checks():
+    return [
+        Check(
+            domain="example.com",
+            response_time=51,
+            status=200,
+            url="https://example.com",
+        ),
+        Check(
+            domain="example.com/200",
+            response_time=65,
+            status=200,
+            url="https://example.com",
+        ),
+        Check(
+            domain="example.com/404",
+            response_time=35,
+            status=404,
+            url="https://example.com",
+        ),
+    ]

tests/test_consumer.py

@@ -1,9 +1,7 @@
 import asyncio
-
-import aiokafka
-from mock import AsyncMock, Mock, patch
 import pytest
 
+from mock import AsyncMock, Mock, patch
 from chweb.consumer import Consumer
 
 
@@ -13,10 +11,9 @@ 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()
+    consumer.consumer.__aiter__.assert_called()
     task.cancel()

tests/test_db.py

@@ -0,0 +1,37 @@
+import asyncio
+
+import pytest
+from mock import AsyncMock, Mock
+
+import chweb.consumer
+from chweb.consumer import DbWriter
+
+
+@pytest.mark.asyncio
+async def test_db_setup(three_checks, config, event_loop):
+    queue = asyncio.Queue()
+    db_writer = DbWriter(config, Mock(), event_loop, queue)
+    assert db_writer.conn is None
+
+    chweb.consumer.asyncpg = AsyncMock()
+    await db_writer.connect()
+    chweb.consumer.asyncpg.connect.assert_called()
+
+    db_writer.conn = AsyncMock()
+    await db_writer.setup()
+    db_writer.conn.execute.assert_called()
+
+    for check in three_checks:
+        await queue.put(check)
+    print("&&&&&&&&&&")
+    print("&&&&&&&&&&")
+    print("&&&&&&&&&&")
+    print(queue.qsize())
+    print(config.postgres.batch_size)
+    print("&&&&&&&&&&")
+    print("&&&&&&&&&&")
+    print("&&&&&&&&&&")
+    task = event_loop.create_task(db_writer())
+    await asyncio.sleep(0.5)
+    db_writer.conn.copy_records_to_table.assert_called()
+    task.cancel()