Fix flaky test, update ty, review linter

nucliadb/src/nucliadb/search/api/v1/augment.py

@@ -240,7 +240,7 @@ def parse_first_augments(item: AugmentRequest) -> list[Augment]:
             if len(file_select) > 0:
                 augmentations.append(
                     FileAugment(
-                        given=given,  # type: ignore
+                        given=given,  # type: ignore[arg-type]
                         select=file_select,
                     )
                 )
@@ -277,7 +277,7 @@ def parse_first_augments(item: AugmentRequest) -> list[Augment]:
             if len(conversation_select) > 0:
                 augmentations.append(
                     ConversationAugment(
-                        given=given,  # type: ignore
+                        given=given,  # type: ignore[arg-type]
                         select=conversation_select,
                     )
                 )

nucliadb/src/nucliadb/search/augmentor/augmentor.py

@@ -113,27 +113,21 @@ def __init__(self, kbid: str):
     async def parse(self, augmentations: list[Augment]):
         for augmentation in augmentations:
             if augmentation.from_ == "resources":
-                augmentation = cast(ResourceAugment, augmentation)
                 self._parse_resource(augmentation)
 
             elif augmentation.from_ == "resources.deep":
-                augmentation = cast(DeepResourceAugment, augmentation)
                 self._parse_deep_resource(augmentation)
 
             elif augmentation.from_ == "fields":
-                augmentation = cast(FieldAugment, augmentation)
                 await self._parse_field(augmentation)
 
             elif augmentation.from_ == "files":
-                augmentation = cast(FileAugment, augmentation)
                 self._parse_file_field(augmentation)
 
             elif augmentation.from_ == "conversations":
-                augmentation = cast(ConversationAugment, augmentation)
                 self._parse_conversation_field(augmentation)
 
             elif augmentation.from_ == "paragraphs":
-                augmentation = cast(ParagraphAugment, augmentation)
                 self._parse_paragraph(augmentation)
 
             else:  # pragma: no cover

nucliadb/src/nucliadb/search/predict.py

@@ -604,7 +604,7 @@ async def run_agents(self, kbid: str, item: RunAgentsRequest) -> RunAgentsRespon
         ada = AppliedDataAugmentation()
         serialized_fm = base64.b64encode(fm.SerializeToString()).decode("utf-8")
         augmented_field = AugmentedField(
-            metadata=serialized_fm,  # type: ignore
+            metadata=serialized_fm,  # type:ignore[arg-type]
             applied_data_augmentation=ada,
             input_nuclia_tokens=1.0,
             output_nuclia_tokens=1.0,

nucliadb/src/nucliadb/search/search/find.py

@@ -112,7 +112,7 @@ async def _ndb_index_find(
         )
         text_block_hydration_options = TextBlockHydrationOptions(
             highlight=item.highlight,
-            ematches=pb_response.paragraph.ematches,  # type: ignore
+            ematches=list(pb_response.paragraph.ematches),
         )
         search_results = await build_find_response(
             pb_response,

nucliadb/tests/ingest/unit/orm/test_resource.py

@@ -212,22 +212,22 @@ async def test_get_fields_ids_caches_keys(txn, storage, kb: str, rid: str):
         (FieldType.TEXT, "bar"),
     ]
     new_field_keys: list[tuple[FieldType.ValueType, str]] = [(FieldType.LINK, "baz")]
-    resource._inner_get_fields_ids = AsyncMock(return_value=new_field_keys)  # type: ignore
+    resource._inner_get_fields_ids = AsyncMock(return_value=new_field_keys)  # type: ignore[method-assign]
     resource.all_fields_keys = cached_field_keys
 
     assert await resource.get_fields_ids() == cached_field_keys
-    resource._inner_get_fields_ids.assert_not_awaited()  # type: ignore[ty:unresolved-attribute]
+    resource._inner_get_fields_ids.assert_not_awaited()
 
     assert await resource.get_fields_ids(force=True) == new_field_keys
-    resource._inner_get_fields_ids.assert_awaited_once()  # type: ignore[ty:unresolved-attribute]
+    resource._inner_get_fields_ids.assert_awaited_once()
     assert resource.all_fields_keys == new_field_keys
 
     # If the all_field_keys is an empty list,
     # we should not be calling the inner_get_fields_ids
     resource.all_fields_keys = []
-    resource._inner_get_fields_ids.reset_mock()  # type: ignore[ty:unresolved-attribute]
+    resource._inner_get_fields_ids.reset_mock()
     assert await resource.get_fields_ids() == []
-    resource._inner_get_fields_ids.assert_not_awaited()  # type: ignore[ty:unresolved-attribute]
+    resource._inner_get_fields_ids.assert_not_awaited()
 
 
 async def test_get_set_all_field_ids(txn, storage, kb, rid):
@@ -313,27 +313,29 @@ async def test_update_all_fields_key(txn, storage, kb, rid):
 
 async def test_apply_fields_calls_update_all_field_ids(txn, storage, kb, rid):
     resource = Resource(txn, storage, kb, rid)
-    resource.update_all_field_ids = AsyncMock()  # type: ignore
-    resource.set_field = AsyncMock()  # type: ignore
+    resource.update_all_field_ids = AsyncMock()
+    resource.set_field = AsyncMock()
+    resource.delete_field = AsyncMock()
+    resource.set_file_field_md5 = AsyncMock()
 
     bm = MagicMock()
     bm.texts = {"text": MagicMock()}
     bm.links = {"link": MagicMock()}
     bm.files = {"file": MagicMock()}
     bm.conversations = {"conversation": MagicMock()}
-    bm.delete_fields.append(FieldID(field_type=FieldType.CONVERSATION, field="to_delete"))
+    bm.delete_fields = [FieldID(field_type=FieldType.CONVERSATION, field="to_delete")]
 
-    await resource.apply_field_values(bm)  # ty:ignore[invalid-argument-type]
+    await resource.apply_field_values(bm)
 
-    resource.update_all_field_ids.assert_awaited_once()  # type: ignore[ty:unresolved-attribute]
+    resource.update_all_field_ids.assert_awaited_once()
 
-    resource.update_all_field_ids.call_args[1]["updated"] == [  # ty:ignore[unresolved-attribute]
+    assert resource.update_all_field_ids.call_args[1]["updated"] == [
         FieldID(field_type=FieldType.TEXT, field="text"),
         FieldID(field_type=FieldType.LINK, field="link"),
         FieldID(field_type=FieldType.FILE, field="file"),
         FieldID(field_type=FieldType.CONVERSATION, field="conversation"),
     ]
-    resource.update_all_field_ids.call_args[1]["deleted"] == [  # ty:ignore[unresolved-attribute]
+    assert resource.update_all_field_ids.call_args[1]["deleted"] == [
         FieldID(field_type=FieldType.CONVERSATION, field="to_delete"),
     ]
 

nucliadb/tests/ingest/unit/service/test_writer.py

@@ -57,7 +57,7 @@ def field_value(self):
     def field(self, field_value):
         val = Text("id", Mock())
         val.value = field_value.SerializeToString()
-        val.set_vectors = AsyncMock()  # ty:ignore[invalid-assignment]
+        val.set_vectors = AsyncMock()
         yield val
 
     @pytest.fixture(scope="function")

nucliadb/tests/nucliadb/ask/test_ask_security.py

@@ -20,6 +20,7 @@
 
 import asyncio
 import json
+from typing import Any
 
 import pytest
 from httpx import AsyncClient
@@ -227,12 +228,12 @@ async def _test_ask_request_with_security(
     security_groups: list[str] | None,
     expected_resources: list[str],
 ):
-    payload = {
+    payload: dict[str, Any] = {
         "query": query,
     }
     headers = {"x_synchronous": "true"}
     if security_groups:
-        payload["security"] = {"groups": security_groups}  # type: ignore
+        payload["security"] = {"groups": security_groups}
 
     if ask_endpoint == "ask_post":
         resp = await nucliadb_reader.post(f"/kb/{kbid}/ask", json=payload, headers=headers)

nucliadb/tests/nucliadb/integration/search/test_search.py

@@ -225,7 +225,6 @@ def broker_resource_with_classifications(standalone_knowledgebox):
     c1.label = "label1"
     c1.labelset = "labelset1"
     fcm.metadata.metadata.classifications.append(c1)
-    bm.field_metadata.append(fcm)
 
     c2 = rpb.Classification()
     c2.label = "label2"
@@ -252,10 +251,6 @@ def broker_resource_with_classifications(standalone_knowledgebox):
     return bm
 
 
-# FIXME: for some reason, paragraph search here doesn't return any result
-# sometimes, although the resource seems to be properly indexed and synced in
-# nidx. Further debugging is needed
-@pytest.mark.flaky(reruns=5)
 @pytest.mark.deploy_modes("standalone")
 async def test_search_returns_labels(
     nucliadb_search: AsyncClient, nucliadb_ingest_grpc: WriterStub, standalone_knowledgebox
@@ -268,6 +263,7 @@ async def test_search_returns_labels(
     )
     assert resp.status_code == 200
     content = resp.json()
+
     assert content["paragraphs"]["results"]
     par = content["paragraphs"]["results"][0]
     assert set(par["labels"]) == {"labelset1/label2", "labelset1/label1"}

nucliadb/tests/nucliadb/integration/test_security.py

@@ -18,6 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
+from typing import Any
 
 import pytest
 from httpx import AsyncClient
@@ -246,15 +247,15 @@ async def _test_search_request_with_security(
     security_groups: list[str] | None,
     expected_resources: list[str],
 ):
-    payload = {
+    payload: dict[str, Any] = {
         "query": query,
     }
-    params = {
+    params: dict[str, Any] = {
         "query": query,
     }
     if security_groups is not None:
-        payload["security"] = {"groups": security_groups}  # type: ignore
-        params["security_groups"] = security_groups  # type: ignore
+        payload["security"] = {"groups": security_groups}
+        params["security_groups"] = security_groups
 
     if method == "POST" and endpoint == "find":
         resp = await nucliadb_reader.post(

nucliadb/tests/nucliadb/integration/test_vectorsets.py

@@ -308,6 +308,7 @@ async def inner(*args, **kwargs):
         node_search_spy, result, error = query
         assert result is not None
         assert error is None
+        assert node_search_spy is not None
 
         request = node_search_spy.call_args[0][0]
         # there's only one model and we get it as the default
@@ -326,6 +327,7 @@ async def inner(*args, **kwargs):
         node_search_spy, result, error = query
         assert result is not None
         assert error is None
+        assert node_search_spy is not None
 
         request = node_search_spy.call_args[0][0]
         assert request.vectorset == "model"
@@ -375,6 +377,7 @@ async def inner(*args, **kwargs):
             node_search_spy, result, error = query
             assert result is not None
             assert error is None
+            assert node_search_spy is not None
 
             request = node_search_spy.call_args[0][0]
             assert request.vectorset == "model-A"
@@ -399,6 +402,7 @@ async def inner(*args, **kwargs):
             node_search_spy, result, error = query
             assert result is not None
             assert error is None
+            assert node_search_spy is not None
 
             request = node_search_spy.call_args[0][0]
             assert request.vectorset == "model-B"
@@ -421,6 +425,8 @@ async def inner(*args, **kwargs):
             )
             assert resp.status_code == 200
             node_search_spy, result, error = query
+            assert node_search_spy is not None
+
             request = node_search_spy.call_args[0][0]
             assert result is not None
             assert error is None

nucliadb/tests/nucliadb/unit/common/test_back_pressure.py

@@ -140,28 +140,28 @@ async def test_check_processing_behind(settings, cache, nats_conn):
     settings.max_processing_pending = 5
 
     materializer = BackPressureMaterializer(nats_conn)
-    materializer.get_processing_pending = mock.AsyncMock(return_value=1)  # ty:ignore[invalid-assignment]
+    materializer.get_processing_pending = mock.AsyncMock(return_value=1)
 
     # Check that it runs and does not raise an exception if the pending is low
     await materializer.check_processing("kbid")
-    materializer.get_processing_pending.assert_awaited_once_with("kbid")  # ty:ignore[unresolved-attribute]
+    materializer.get_processing_pending.assert_awaited_once_with("kbid")
 
     # Check that it raises an exception if the pending is too high
-    materializer.get_processing_pending.reset_mock()  # ty:ignore[unresolved-attribute]
-    materializer.get_processing_pending.return_value = 10  # ty:ignore[unresolved-attribute]
+    materializer.get_processing_pending.reset_mock()
+    materializer.get_processing_pending.return_value = 10
     with pytest.raises(BackPressureException):
         await materializer.check_processing("kbid")
-    materializer.get_processing_pending.assert_awaited_once_with("kbid")  # ty:ignore[unresolved-attribute]
+    materializer.get_processing_pending.assert_awaited_once_with("kbid")
 
 
 async def test_check_processing_behind_does_not_run_if_configured_max_is_zero(settings, cache):
     settings.max_processing_pending = 0
     materializer = BackPressureMaterializer(mock.Mock())
-    materializer.get_processing_pending = mock.AsyncMock(return_value=100)  # ty:ignore[invalid-assignment]
+    materializer.get_processing_pending = mock.AsyncMock(return_value=100)
 
     await materializer.check_processing("kbid")
 
-    materializer.get_processing_pending.assert_not_called()  # ty:ignore[unresolved-attribute]
+    materializer.get_processing_pending.assert_not_called()
 
 
 async def test_check_ingest_behind(settings, cache):

nucliadb/tests/nucliadb/unit/common/test_context.py

@@ -25,12 +25,12 @@
 
 async def test_initialize_happens_only_once():
     context = ApplicationContext()
-    context._initialize = AsyncMock()  # ty:ignore[invalid-assignment]
+    context._initialize = AsyncMock()
 
     tasks = []
     for _ in range(10):
         tasks.append(context.initialize())
     await asyncio.gather(*tasks)
 
-    context._initialize.assert_awaited_once()  # ty:ignore[unresolved-attribute]
-    context._initialized is True
+    context._initialize.assert_awaited_once()
+    assert context._initialized is True

nucliadb/tests/nucliadb/unit/export_import/test_datamanager.py

@@ -27,11 +27,11 @@ async def test_try_delete_from_storage():
     storage = Mock()
 
     dm = ExportImportDataManager(driver, storage)
-    dm.delete_export = AsyncMock(side_effect=ValueError())  # type: ignore
-    dm.delete_import = AsyncMock(side_effect=ValueError())  # type: ignore
+    dm.delete_export = AsyncMock(side_effect=ValueError())
+    dm.delete_import = AsyncMock(side_effect=ValueError())
 
     await dm.try_delete_from_storage("export", "kbid", "foo")
     await dm.try_delete_from_storage("import", "kbid", "foo")
 
-    dm.delete_export.assert_called()  # ty:ignore[unresolved-attribute]
-    dm.delete_import.assert_called()  # ty:ignore[unresolved-attribute]
+    dm.delete_export.assert_called()
+    dm.delete_import.assert_called()

nucliadb/tests/search/integration/requesters/test_utils.py

@@ -44,8 +44,6 @@ async def test_vector_result_metadata(nucliadb_search: AsyncClient, test_search_
         SearchRequest(
             query="own text",
             features=[SearchOptions.SEMANTIC],
-            label_filters=[],  # type: ignore
-            keyword_filters=[],  # type: ignore
             faceted=[],
             top_k=20,
             min_score=MinScore(bm25=0, semantic=-1),

nucliadb/tests/search/unit/search/search/test_fetcher_cache.py

@@ -67,5 +67,5 @@ def new_fetcher() -> Fetcher:
         generative_model=None,
         query_image=None,
     )
-    fetcher.get_vectorset = AsyncMock(return_value=vectorset)  # type: ignore
+    fetcher.get_vectorset = AsyncMock(return_value=vectorset)  # type: ignore[method-assign]
     return fetcher

nucliadb/tests/train/test_get_entities.py

@@ -34,7 +34,7 @@ async def entities_manager_mock():
     original = nodes.EntitiesManager
 
     mock = Mock()
-    nodes.EntitiesManager = Mock(return_value=mock)  # ty:ignore[invalid-assignment]
+    nodes.EntitiesManager = Mock(return_value=mock)
 
     yield mock
 

nucliadb_models/tests/test_metadata.py

@@ -27,7 +27,7 @@ def test_relation_validator():
     )
 
     with pytest.raises(ValidationError):
-        metadata.UserMetadata(relations=["my-wrong-relation"])  # type: ignore
+        metadata.UserMetadata(relations=["my-wrong-relation"])
 
 
 def test_relation_entity_model_validator():

nucliadb_models/tests/test_search.py

@@ -65,13 +65,13 @@ def test_field_extension_strategy_fields_field_validator():
 
 def test_find_request_fulltext_feature_not_allowed():
     with pytest.raises(ValidationError):
-        search.FindRequest(features=[search.SearchOptions.FULLTEXT])  # type: ignore
+        search.FindRequest(features=[search.SearchOptions.FULLTEXT])
 
 
 def test_find_supports_search_options():
-    search.FindRequest(features=[search.SearchOptions.KEYWORD])  # type: ignore
-    search.FindRequest(features=[search.SearchOptions.SEMANTIC])  # type: ignore
-    search.FindRequest(features=[search.SearchOptions.RELATIONS])  # type: ignore
+    search.FindRequest(features=[search.SearchOptions.KEYWORD])
+    search.FindRequest(features=[search.SearchOptions.SEMANTIC])
+    search.FindRequest(features=[search.SearchOptions.RELATIONS])
 
 
 def test_search_semantic_with_offset_not_supported():
@@ -149,13 +149,13 @@ def test_rank_fusion(rank_fusion, expected):
 
 def test_rank_fusion_errors():
     with pytest.raises(ValueError):
-        search.FindRequest(rank_fusion="unknown")  # type: ignore
+        search.FindRequest(rank_fusion="unknown")
     with pytest.raises(ValueError):
-        search.AskRequest(query="q", rank_fusion="unknown")  # type: ignore
+        search.AskRequest(query="q", rank_fusion="unknown")
 
 
 def test_legacy_rank_fusion_fix():
-    req = search.FindRequest(rank_fusion="legacy")  # type: ignore
+    req = search.FindRequest(rank_fusion="legacy")
     assert req.rank_fusion == "rrf"
 
     req = search.FindRequest.model_validate({"rank_fusion": "legacy"})