Fetch operations before starting WS connection (#7)

Author: droserasprout

src/dipdup/datasources/tzkt/datasource.py

@@ -72,8 +72,15 @@ def _get_client(self) -> BaseHubConnection:
 
     async def start(self):
         self._logger.info('Starting datasource')
-        for config in self._operation_index_configs.values():
-            await self.add_subscription(config.contract_config.address)
+        for operation_index_config in self._operation_index_configs.values():
+            await self.add_subscription(operation_index_config.contract)
+
+            latest_block = await self.get_latest_block()
+            current_level = latest_block['level']
+            state_level = operation_index_config.state.level
+
+            if current_level != state_level:
+                await self.fetch_operations(state_level)
 
         self._logger.info('Starting websocket client')
         await self._get_client().start()
@@ -299,3 +306,17 @@ def convert_operation(cls, operation_json: Dict[str, Any]) -> OperationData:
             initiator_address=operation_json.get('initiator', {}).get('address'),
             parameter=operation_json.get('parameters'),
         )
+
+    async def get_latest_block(self) -> Dict[str, Any]:
+        self._logger.info('Fetching latest block')
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                url=f'{self._url}/v1/blocks',
+                params={
+                    "limit": 1,
+                    "sort.desc": "id",
+                },
+            ) as resp:
+                blocks = await resp.json()
+        self._logger.debug(blocks)
+        return blocks[0]

tests/test_dipdup/test_datasources/test_tzkt/test_datasource.py

@@ -57,9 +57,17 @@ async def test_start(self):
         client = self.datasource._get_client()
         client.start = AsyncMock()
 
-        await self.datasource.start()
+        fetch_operations_mock = AsyncMock()
+        self.datasource.fetch_operations = fetch_operations_mock
+
+        get_mock = MagicMock()
+        get_mock.return_value.__aenter__.return_value.json.return_value = [{'level': 1337}]
+
+        with patch('aiohttp.ClientSession.get', get_mock):
+            await self.datasource.start()
 
-        self.assertEqual({self.index_config.contract.address: ['transaction']}, self.datasource._subscriptions)
+        fetch_operations_mock.assert_awaited_with(0)
+        self.assertEqual({self.index_config.contract: ['transaction']}, self.datasource._subscriptions)
         client.start.assert_awaited()
 
     async def test_on_connect_subscribe_to_operations(self):