nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (418) hide show
  1. migrations/0003_allfields_key.py +1 -35
  2. migrations/0009_upgrade_relations_and_texts_to_v2.py +4 -2
  3. migrations/0010_fix_corrupt_indexes.py +10 -10
  4. migrations/0011_materialize_labelset_ids.py +1 -16
  5. migrations/0012_rollover_shards.py +5 -10
  6. migrations/0014_rollover_shards.py +4 -5
  7. migrations/0015_targeted_rollover.py +5 -10
  8. migrations/0016_upgrade_to_paragraphs_v2.py +25 -28
  9. migrations/0017_multiple_writable_shards.py +2 -4
  10. migrations/0018_purge_orphan_kbslugs.py +5 -7
  11. migrations/0019_upgrade_to_paragraphs_v3.py +25 -28
  12. migrations/0020_drain_nodes_from_cluster.py +3 -3
  13. nucliadb/standalone/tests/unit/test_run.py → migrations/0021_overwrite_vectorsets_key.py +16 -19
  14. nucliadb/tests/unit/test_openapi.py → migrations/0022_fix_paragraph_deletion_bug.py +16 -11
  15. migrations/0023_backfill_pg_catalog.py +80 -0
  16. migrations/0025_assign_models_to_kbs_v2.py +113 -0
  17. migrations/0026_fix_high_cardinality_content_types.py +61 -0
  18. migrations/0027_rollover_texts3.py +73 -0
  19. nucliadb/ingest/fields/date.py → migrations/pg/0001_bootstrap.py +10 -12
  20. migrations/pg/0002_catalog.py +42 -0
  21. nucliadb/ingest/tests/unit/test_settings.py → migrations/pg/0003_catalog_kbid_index.py +5 -3
  22. nucliadb/common/cluster/base.py +30 -16
  23. nucliadb/common/cluster/discovery/base.py +6 -14
  24. nucliadb/common/cluster/discovery/k8s.py +9 -19
  25. nucliadb/common/cluster/discovery/manual.py +1 -3
  26. nucliadb/common/cluster/discovery/utils.py +1 -3
  27. nucliadb/common/cluster/grpc_node_dummy.py +3 -11
  28. nucliadb/common/cluster/index_node.py +10 -19
  29. nucliadb/common/cluster/manager.py +174 -59
  30. nucliadb/common/cluster/rebalance.py +27 -29
  31. nucliadb/common/cluster/rollover.py +353 -194
  32. nucliadb/common/cluster/settings.py +6 -0
  33. nucliadb/common/cluster/standalone/grpc_node_binding.py +13 -64
  34. nucliadb/common/cluster/standalone/index_node.py +4 -11
  35. nucliadb/common/cluster/standalone/service.py +2 -6
  36. nucliadb/common/cluster/standalone/utils.py +2 -6
  37. nucliadb/common/cluster/utils.py +29 -22
  38. nucliadb/common/constants.py +20 -0
  39. nucliadb/common/context/__init__.py +3 -0
  40. nucliadb/common/context/fastapi.py +8 -5
  41. nucliadb/{tests/knowledgeboxes/__init__.py → common/counters.py} +8 -2
  42. nucliadb/common/datamanagers/__init__.py +7 -1
  43. nucliadb/common/datamanagers/atomic.py +22 -4
  44. nucliadb/common/datamanagers/cluster.py +5 -5
  45. nucliadb/common/datamanagers/entities.py +6 -16
  46. nucliadb/common/datamanagers/fields.py +84 -0
  47. nucliadb/common/datamanagers/kb.py +83 -37
  48. nucliadb/common/datamanagers/labels.py +26 -56
  49. nucliadb/common/datamanagers/processing.py +2 -6
  50. nucliadb/common/datamanagers/resources.py +41 -103
  51. nucliadb/common/datamanagers/rollover.py +76 -15
  52. nucliadb/common/datamanagers/synonyms.py +1 -1
  53. nucliadb/common/datamanagers/utils.py +15 -6
  54. nucliadb/common/datamanagers/vectorsets.py +110 -0
  55. nucliadb/common/external_index_providers/base.py +257 -0
  56. nucliadb/{ingest/tests/unit/orm/test_orm_utils.py → common/external_index_providers/exceptions.py} +9 -8
  57. nucliadb/common/external_index_providers/manager.py +101 -0
  58. nucliadb/common/external_index_providers/pinecone.py +933 -0
  59. nucliadb/common/external_index_providers/settings.py +52 -0
  60. nucliadb/common/http_clients/auth.py +3 -6
  61. nucliadb/common/http_clients/processing.py +6 -11
  62. nucliadb/common/http_clients/utils.py +1 -3
  63. nucliadb/common/ids.py +240 -0
  64. nucliadb/common/locking.py +29 -7
  65. nucliadb/common/maindb/driver.py +11 -35
  66. nucliadb/common/maindb/exceptions.py +3 -0
  67. nucliadb/common/maindb/local.py +22 -9
  68. nucliadb/common/maindb/pg.py +206 -111
  69. nucliadb/common/maindb/utils.py +11 -42
  70. nucliadb/common/models_utils/from_proto.py +479 -0
  71. nucliadb/common/models_utils/to_proto.py +60 -0
  72. nucliadb/common/nidx.py +260 -0
  73. nucliadb/export_import/datamanager.py +25 -19
  74. nucliadb/export_import/exporter.py +5 -11
  75. nucliadb/export_import/importer.py +5 -7
  76. nucliadb/export_import/models.py +3 -3
  77. nucliadb/export_import/tasks.py +4 -4
  78. nucliadb/export_import/utils.py +25 -37
  79. nucliadb/health.py +1 -3
  80. nucliadb/ingest/app.py +15 -11
  81. nucliadb/ingest/consumer/auditing.py +21 -19
  82. nucliadb/ingest/consumer/consumer.py +82 -47
  83. nucliadb/ingest/consumer/materializer.py +5 -12
  84. nucliadb/ingest/consumer/pull.py +12 -27
  85. nucliadb/ingest/consumer/service.py +19 -17
  86. nucliadb/ingest/consumer/shard_creator.py +2 -4
  87. nucliadb/ingest/consumer/utils.py +1 -3
  88. nucliadb/ingest/fields/base.py +137 -105
  89. nucliadb/ingest/fields/conversation.py +18 -5
  90. nucliadb/ingest/fields/exceptions.py +1 -4
  91. nucliadb/ingest/fields/file.py +7 -16
  92. nucliadb/ingest/fields/link.py +5 -10
  93. nucliadb/ingest/fields/text.py +9 -4
  94. nucliadb/ingest/orm/brain.py +200 -213
  95. nucliadb/ingest/orm/broker_message.py +181 -0
  96. nucliadb/ingest/orm/entities.py +36 -51
  97. nucliadb/ingest/orm/exceptions.py +12 -0
  98. nucliadb/ingest/orm/knowledgebox.py +322 -197
  99. nucliadb/ingest/orm/processor/__init__.py +2 -700
  100. nucliadb/ingest/orm/processor/auditing.py +4 -23
  101. nucliadb/ingest/orm/processor/data_augmentation.py +164 -0
  102. nucliadb/ingest/orm/processor/pgcatalog.py +84 -0
  103. nucliadb/ingest/orm/processor/processor.py +752 -0
  104. nucliadb/ingest/orm/processor/sequence_manager.py +1 -1
  105. nucliadb/ingest/orm/resource.py +249 -403
  106. nucliadb/ingest/orm/utils.py +4 -4
  107. nucliadb/ingest/partitions.py +3 -9
  108. nucliadb/ingest/processing.py +70 -73
  109. nucliadb/ingest/py.typed +0 -0
  110. nucliadb/ingest/serialize.py +37 -167
  111. nucliadb/ingest/service/__init__.py +1 -3
  112. nucliadb/ingest/service/writer.py +185 -412
  113. nucliadb/ingest/settings.py +10 -20
  114. nucliadb/ingest/utils.py +3 -6
  115. nucliadb/learning_proxy.py +242 -55
  116. nucliadb/metrics_exporter.py +30 -19
  117. nucliadb/middleware/__init__.py +1 -3
  118. nucliadb/migrator/command.py +1 -3
  119. nucliadb/migrator/datamanager.py +13 -13
  120. nucliadb/migrator/migrator.py +47 -30
  121. nucliadb/migrator/utils.py +18 -10
  122. nucliadb/purge/__init__.py +139 -33
  123. nucliadb/purge/orphan_shards.py +7 -13
  124. nucliadb/reader/__init__.py +1 -3
  125. nucliadb/reader/api/models.py +1 -12
  126. nucliadb/reader/api/v1/__init__.py +0 -1
  127. nucliadb/reader/api/v1/download.py +21 -88
  128. nucliadb/reader/api/v1/export_import.py +1 -1
  129. nucliadb/reader/api/v1/knowledgebox.py +10 -10
  130. nucliadb/reader/api/v1/learning_config.py +2 -6
  131. nucliadb/reader/api/v1/resource.py +62 -88
  132. nucliadb/reader/api/v1/services.py +64 -83
  133. nucliadb/reader/app.py +12 -29
  134. nucliadb/reader/lifecycle.py +18 -4
  135. nucliadb/reader/py.typed +0 -0
  136. nucliadb/reader/reader/notifications.py +10 -28
  137. nucliadb/search/__init__.py +1 -3
  138. nucliadb/search/api/v1/__init__.py +1 -2
  139. nucliadb/search/api/v1/ask.py +17 -10
  140. nucliadb/search/api/v1/catalog.py +184 -0
  141. nucliadb/search/api/v1/feedback.py +16 -24
  142. nucliadb/search/api/v1/find.py +36 -36
  143. nucliadb/search/api/v1/knowledgebox.py +89 -60
  144. nucliadb/search/api/v1/resource/ask.py +2 -8
  145. nucliadb/search/api/v1/resource/search.py +49 -70
  146. nucliadb/search/api/v1/search.py +44 -210
  147. nucliadb/search/api/v1/suggest.py +39 -54
  148. nucliadb/search/app.py +12 -32
  149. nucliadb/search/lifecycle.py +10 -3
  150. nucliadb/search/predict.py +136 -187
  151. nucliadb/search/py.typed +0 -0
  152. nucliadb/search/requesters/utils.py +25 -58
  153. nucliadb/search/search/cache.py +149 -20
  154. nucliadb/search/search/chat/ask.py +571 -123
  155. nucliadb/search/{tests/unit/test_run.py → search/chat/exceptions.py} +14 -14
  156. nucliadb/search/search/chat/images.py +41 -17
  157. nucliadb/search/search/chat/prompt.py +817 -266
  158. nucliadb/search/search/chat/query.py +213 -309
  159. nucliadb/{tests/migrations/__init__.py → search/search/cut.py} +8 -8
  160. nucliadb/search/search/fetch.py +43 -36
  161. nucliadb/search/search/filters.py +9 -15
  162. nucliadb/search/search/find.py +214 -53
  163. nucliadb/search/search/find_merge.py +408 -391
  164. nucliadb/search/search/hydrator.py +191 -0
  165. nucliadb/search/search/merge.py +187 -223
  166. nucliadb/search/search/metrics.py +73 -2
  167. nucliadb/search/search/paragraphs.py +64 -106
  168. nucliadb/search/search/pgcatalog.py +233 -0
  169. nucliadb/search/search/predict_proxy.py +1 -1
  170. nucliadb/search/search/query.py +305 -150
  171. nucliadb/search/search/query_parser/exceptions.py +22 -0
  172. nucliadb/search/search/query_parser/models.py +101 -0
  173. nucliadb/search/search/query_parser/parser.py +183 -0
  174. nucliadb/search/search/rank_fusion.py +204 -0
  175. nucliadb/search/search/rerankers.py +270 -0
  176. nucliadb/search/search/shards.py +3 -32
  177. nucliadb/search/search/summarize.py +7 -18
  178. nucliadb/search/search/utils.py +27 -4
  179. nucliadb/search/settings.py +15 -1
  180. nucliadb/standalone/api_router.py +4 -10
  181. nucliadb/standalone/app.py +8 -14
  182. nucliadb/standalone/auth.py +7 -21
  183. nucliadb/standalone/config.py +7 -10
  184. nucliadb/standalone/lifecycle.py +26 -25
  185. nucliadb/standalone/migrations.py +1 -3
  186. nucliadb/standalone/purge.py +1 -1
  187. nucliadb/standalone/py.typed +0 -0
  188. nucliadb/standalone/run.py +3 -6
  189. nucliadb/standalone/settings.py +9 -16
  190. nucliadb/standalone/versions.py +15 -5
  191. nucliadb/tasks/consumer.py +8 -12
  192. nucliadb/tasks/producer.py +7 -6
  193. nucliadb/tests/config.py +53 -0
  194. nucliadb/train/__init__.py +1 -3
  195. nucliadb/train/api/utils.py +1 -2
  196. nucliadb/train/api/v1/shards.py +1 -1
  197. nucliadb/train/api/v1/trainset.py +2 -4
  198. nucliadb/train/app.py +10 -31
  199. nucliadb/train/generator.py +10 -19
  200. nucliadb/train/generators/field_classifier.py +7 -19
  201. nucliadb/train/generators/field_streaming.py +156 -0
  202. nucliadb/train/generators/image_classifier.py +12 -18
  203. nucliadb/train/generators/paragraph_classifier.py +5 -9
  204. nucliadb/train/generators/paragraph_streaming.py +6 -9
  205. nucliadb/train/generators/question_answer_streaming.py +19 -20
  206. nucliadb/train/generators/sentence_classifier.py +9 -15
  207. nucliadb/train/generators/token_classifier.py +48 -39
  208. nucliadb/train/generators/utils.py +14 -18
  209. nucliadb/train/lifecycle.py +7 -3
  210. nucliadb/train/nodes.py +23 -32
  211. nucliadb/train/py.typed +0 -0
  212. nucliadb/train/servicer.py +13 -21
  213. nucliadb/train/settings.py +2 -6
  214. nucliadb/train/types.py +13 -10
  215. nucliadb/train/upload.py +3 -6
  216. nucliadb/train/uploader.py +19 -23
  217. nucliadb/train/utils.py +1 -1
  218. nucliadb/writer/__init__.py +1 -3
  219. nucliadb/{ingest/fields/keywordset.py → writer/api/utils.py} +13 -10
  220. nucliadb/writer/api/v1/export_import.py +67 -14
  221. nucliadb/writer/api/v1/field.py +16 -269
  222. nucliadb/writer/api/v1/knowledgebox.py +218 -68
  223. nucliadb/writer/api/v1/resource.py +68 -88
  224. nucliadb/writer/api/v1/services.py +51 -70
  225. nucliadb/writer/api/v1/slug.py +61 -0
  226. nucliadb/writer/api/v1/transaction.py +67 -0
  227. nucliadb/writer/api/v1/upload.py +143 -117
  228. nucliadb/writer/app.py +6 -43
  229. nucliadb/writer/back_pressure.py +16 -38
  230. nucliadb/writer/exceptions.py +0 -4
  231. nucliadb/writer/lifecycle.py +21 -15
  232. nucliadb/writer/py.typed +0 -0
  233. nucliadb/writer/resource/audit.py +2 -1
  234. nucliadb/writer/resource/basic.py +48 -46
  235. nucliadb/writer/resource/field.py +37 -128
  236. nucliadb/writer/resource/origin.py +1 -2
  237. nucliadb/writer/settings.py +6 -2
  238. nucliadb/writer/tus/__init__.py +17 -15
  239. nucliadb/writer/tus/azure.py +111 -0
  240. nucliadb/writer/tus/dm.py +17 -5
  241. nucliadb/writer/tus/exceptions.py +1 -3
  242. nucliadb/writer/tus/gcs.py +49 -84
  243. nucliadb/writer/tus/local.py +21 -37
  244. nucliadb/writer/tus/s3.py +28 -68
  245. nucliadb/writer/tus/storage.py +5 -56
  246. nucliadb/writer/vectorsets.py +125 -0
  247. nucliadb-6.2.1.post2798.dist-info/METADATA +148 -0
  248. nucliadb-6.2.1.post2798.dist-info/RECORD +343 -0
  249. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/WHEEL +1 -1
  250. nucliadb/common/maindb/redis.py +0 -194
  251. nucliadb/common/maindb/tikv.py +0 -433
  252. nucliadb/ingest/fields/layout.py +0 -58
  253. nucliadb/ingest/tests/conftest.py +0 -30
  254. nucliadb/ingest/tests/fixtures.py +0 -764
  255. nucliadb/ingest/tests/integration/consumer/__init__.py +0 -18
  256. nucliadb/ingest/tests/integration/consumer/test_auditing.py +0 -78
  257. nucliadb/ingest/tests/integration/consumer/test_materializer.py +0 -126
  258. nucliadb/ingest/tests/integration/consumer/test_pull.py +0 -144
  259. nucliadb/ingest/tests/integration/consumer/test_service.py +0 -81
  260. nucliadb/ingest/tests/integration/consumer/test_shard_creator.py +0 -68
  261. nucliadb/ingest/tests/integration/ingest/test_ingest.py +0 -684
  262. nucliadb/ingest/tests/integration/ingest/test_processing_engine.py +0 -95
  263. nucliadb/ingest/tests/integration/ingest/test_relations.py +0 -272
  264. nucliadb/ingest/tests/unit/consumer/__init__.py +0 -18
  265. nucliadb/ingest/tests/unit/consumer/test_auditing.py +0 -139
  266. nucliadb/ingest/tests/unit/consumer/test_consumer.py +0 -69
  267. nucliadb/ingest/tests/unit/consumer/test_pull.py +0 -60
  268. nucliadb/ingest/tests/unit/consumer/test_shard_creator.py +0 -140
  269. nucliadb/ingest/tests/unit/consumer/test_utils.py +0 -67
  270. nucliadb/ingest/tests/unit/orm/__init__.py +0 -19
  271. nucliadb/ingest/tests/unit/orm/test_brain.py +0 -247
  272. nucliadb/ingest/tests/unit/orm/test_brain_vectors.py +0 -74
  273. nucliadb/ingest/tests/unit/orm/test_processor.py +0 -131
  274. nucliadb/ingest/tests/unit/orm/test_resource.py +0 -331
  275. nucliadb/ingest/tests/unit/test_cache.py +0 -31
  276. nucliadb/ingest/tests/unit/test_partitions.py +0 -40
  277. nucliadb/ingest/tests/unit/test_processing.py +0 -171
  278. nucliadb/middleware/transaction.py +0 -117
  279. nucliadb/reader/api/v1/learning_collector.py +0 -63
  280. nucliadb/reader/tests/__init__.py +0 -19
  281. nucliadb/reader/tests/conftest.py +0 -31
  282. nucliadb/reader/tests/fixtures.py +0 -136
  283. nucliadb/reader/tests/test_list_resources.py +0 -75
  284. nucliadb/reader/tests/test_reader_file_download.py +0 -273
  285. nucliadb/reader/tests/test_reader_resource.py +0 -353
  286. nucliadb/reader/tests/test_reader_resource_field.py +0 -219
  287. nucliadb/search/api/v1/chat.py +0 -263
  288. nucliadb/search/api/v1/resource/chat.py +0 -174
  289. nucliadb/search/tests/__init__.py +0 -19
  290. nucliadb/search/tests/conftest.py +0 -33
  291. nucliadb/search/tests/fixtures.py +0 -199
  292. nucliadb/search/tests/node.py +0 -466
  293. nucliadb/search/tests/unit/__init__.py +0 -18
  294. nucliadb/search/tests/unit/api/__init__.py +0 -19
  295. nucliadb/search/tests/unit/api/v1/__init__.py +0 -19
  296. nucliadb/search/tests/unit/api/v1/resource/__init__.py +0 -19
  297. nucliadb/search/tests/unit/api/v1/resource/test_chat.py +0 -98
  298. nucliadb/search/tests/unit/api/v1/test_ask.py +0 -120
  299. nucliadb/search/tests/unit/api/v1/test_chat.py +0 -96
  300. nucliadb/search/tests/unit/api/v1/test_predict_proxy.py +0 -98
  301. nucliadb/search/tests/unit/api/v1/test_summarize.py +0 -99
  302. nucliadb/search/tests/unit/search/__init__.py +0 -18
  303. nucliadb/search/tests/unit/search/requesters/__init__.py +0 -18
  304. nucliadb/search/tests/unit/search/requesters/test_utils.py +0 -211
  305. nucliadb/search/tests/unit/search/search/__init__.py +0 -19
  306. nucliadb/search/tests/unit/search/search/test_shards.py +0 -45
  307. nucliadb/search/tests/unit/search/search/test_utils.py +0 -82
  308. nucliadb/search/tests/unit/search/test_chat_prompt.py +0 -270
  309. nucliadb/search/tests/unit/search/test_fetch.py +0 -108
  310. nucliadb/search/tests/unit/search/test_filters.py +0 -125
  311. nucliadb/search/tests/unit/search/test_paragraphs.py +0 -157
  312. nucliadb/search/tests/unit/search/test_predict_proxy.py +0 -106
  313. nucliadb/search/tests/unit/search/test_query.py +0 -153
  314. nucliadb/search/tests/unit/test_app.py +0 -79
  315. nucliadb/search/tests/unit/test_find_merge.py +0 -112
  316. nucliadb/search/tests/unit/test_merge.py +0 -34
  317. nucliadb/search/tests/unit/test_predict.py +0 -525
  318. nucliadb/standalone/tests/__init__.py +0 -19
  319. nucliadb/standalone/tests/conftest.py +0 -33
  320. nucliadb/standalone/tests/fixtures.py +0 -38
  321. nucliadb/standalone/tests/unit/__init__.py +0 -18
  322. nucliadb/standalone/tests/unit/test_api_router.py +0 -61
  323. nucliadb/standalone/tests/unit/test_auth.py +0 -169
  324. nucliadb/standalone/tests/unit/test_introspect.py +0 -35
  325. nucliadb/standalone/tests/unit/test_migrations.py +0 -63
  326. nucliadb/standalone/tests/unit/test_versions.py +0 -68
  327. nucliadb/tests/benchmarks/__init__.py +0 -19
  328. nucliadb/tests/benchmarks/test_search.py +0 -99
  329. nucliadb/tests/conftest.py +0 -32
  330. nucliadb/tests/fixtures.py +0 -735
  331. nucliadb/tests/knowledgeboxes/philosophy_books.py +0 -202
  332. nucliadb/tests/knowledgeboxes/ten_dummy_resources.py +0 -107
  333. nucliadb/tests/migrations/test_migration_0017.py +0 -76
  334. nucliadb/tests/migrations/test_migration_0018.py +0 -95
  335. nucliadb/tests/tikv.py +0 -240
  336. nucliadb/tests/unit/__init__.py +0 -19
  337. nucliadb/tests/unit/common/__init__.py +0 -19
  338. nucliadb/tests/unit/common/cluster/__init__.py +0 -19
  339. nucliadb/tests/unit/common/cluster/discovery/__init__.py +0 -19
  340. nucliadb/tests/unit/common/cluster/discovery/test_k8s.py +0 -172
  341. nucliadb/tests/unit/common/cluster/standalone/__init__.py +0 -18
  342. nucliadb/tests/unit/common/cluster/standalone/test_service.py +0 -114
  343. nucliadb/tests/unit/common/cluster/standalone/test_utils.py +0 -61
  344. nucliadb/tests/unit/common/cluster/test_cluster.py +0 -408
  345. nucliadb/tests/unit/common/cluster/test_kb_shard_manager.py +0 -173
  346. nucliadb/tests/unit/common/cluster/test_rebalance.py +0 -38
  347. nucliadb/tests/unit/common/cluster/test_rollover.py +0 -282
  348. nucliadb/tests/unit/common/maindb/__init__.py +0 -18
  349. nucliadb/tests/unit/common/maindb/test_driver.py +0 -127
  350. nucliadb/tests/unit/common/maindb/test_tikv.py +0 -53
  351. nucliadb/tests/unit/common/maindb/test_utils.py +0 -92
  352. nucliadb/tests/unit/common/test_context.py +0 -36
  353. nucliadb/tests/unit/export_import/__init__.py +0 -19
  354. nucliadb/tests/unit/export_import/test_datamanager.py +0 -37
  355. nucliadb/tests/unit/export_import/test_utils.py +0 -301
  356. nucliadb/tests/unit/migrator/__init__.py +0 -19
  357. nucliadb/tests/unit/migrator/test_migrator.py +0 -87
  358. nucliadb/tests/unit/tasks/__init__.py +0 -19
  359. nucliadb/tests/unit/tasks/conftest.py +0 -42
  360. nucliadb/tests/unit/tasks/test_consumer.py +0 -92
  361. nucliadb/tests/unit/tasks/test_producer.py +0 -95
  362. nucliadb/tests/unit/tasks/test_tasks.py +0 -58
  363. nucliadb/tests/unit/test_field_ids.py +0 -49
  364. nucliadb/tests/unit/test_health.py +0 -86
  365. nucliadb/tests/unit/test_kb_slugs.py +0 -54
  366. nucliadb/tests/unit/test_learning_proxy.py +0 -252
  367. nucliadb/tests/unit/test_metrics_exporter.py +0 -77
  368. nucliadb/tests/unit/test_purge.py +0 -136
  369. nucliadb/tests/utils/__init__.py +0 -74
  370. nucliadb/tests/utils/aiohttp_session.py +0 -44
  371. nucliadb/tests/utils/broker_messages/__init__.py +0 -171
  372. nucliadb/tests/utils/broker_messages/fields.py +0 -197
  373. nucliadb/tests/utils/broker_messages/helpers.py +0 -33
  374. nucliadb/tests/utils/entities.py +0 -78
  375. nucliadb/train/api/v1/check.py +0 -60
  376. nucliadb/train/tests/__init__.py +0 -19
  377. nucliadb/train/tests/conftest.py +0 -29
  378. nucliadb/train/tests/fixtures.py +0 -342
  379. nucliadb/train/tests/test_field_classification.py +0 -122
  380. nucliadb/train/tests/test_get_entities.py +0 -80
  381. nucliadb/train/tests/test_get_info.py +0 -51
  382. nucliadb/train/tests/test_get_ontology.py +0 -34
  383. nucliadb/train/tests/test_get_ontology_count.py +0 -63
  384. nucliadb/train/tests/test_image_classification.py +0 -221
  385. nucliadb/train/tests/test_list_fields.py +0 -39
  386. nucliadb/train/tests/test_list_paragraphs.py +0 -73
  387. nucliadb/train/tests/test_list_resources.py +0 -39
  388. nucliadb/train/tests/test_list_sentences.py +0 -71
  389. nucliadb/train/tests/test_paragraph_classification.py +0 -123
  390. nucliadb/train/tests/test_paragraph_streaming.py +0 -118
  391. nucliadb/train/tests/test_question_answer_streaming.py +0 -239
  392. nucliadb/train/tests/test_sentence_classification.py +0 -143
  393. nucliadb/train/tests/test_token_classification.py +0 -136
  394. nucliadb/train/tests/utils.py +0 -101
  395. nucliadb/writer/layouts/__init__.py +0 -51
  396. nucliadb/writer/layouts/v1.py +0 -59
  397. nucliadb/writer/tests/__init__.py +0 -19
  398. nucliadb/writer/tests/conftest.py +0 -31
  399. nucliadb/writer/tests/fixtures.py +0 -191
  400. nucliadb/writer/tests/test_fields.py +0 -475
  401. nucliadb/writer/tests/test_files.py +0 -740
  402. nucliadb/writer/tests/test_knowledgebox.py +0 -49
  403. nucliadb/writer/tests/test_reprocess_file_field.py +0 -133
  404. nucliadb/writer/tests/test_resources.py +0 -476
  405. nucliadb/writer/tests/test_service.py +0 -137
  406. nucliadb/writer/tests/test_tus.py +0 -203
  407. nucliadb/writer/tests/utils.py +0 -35
  408. nucliadb/writer/tus/pg.py +0 -125
  409. nucliadb-4.0.0.post542.dist-info/METADATA +0 -135
  410. nucliadb-4.0.0.post542.dist-info/RECORD +0 -462
  411. {nucliadb/ingest/tests → migrations/pg}/__init__.py +0 -0
  412. /nucliadb/{ingest/tests/integration → common/external_index_providers}/__init__.py +0 -0
  413. /nucliadb/{ingest/tests/integration/ingest → common/models_utils}/__init__.py +0 -0
  414. /nucliadb/{ingest/tests/unit → search/search/query_parser}/__init__.py +0 -0
  415. /nucliadb/{ingest/tests → tests}/vectors.py +0 -0
  416. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/entry_points.txt +0 -0
  417. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/top_level.txt +0 -0
  418. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/zip-safe +0 -0
@@ -1,282 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
19
- #
20
- from datetime import datetime
21
- from unittest.mock import ANY, AsyncMock, MagicMock, patch
22
-
23
- import pytest
24
-
25
- from nucliadb.common.cluster import manager, rollover
26
- from nucliadb.common.cluster.index_node import IndexNode
27
- from nucliadb_protos import writer_pb2
28
-
29
- pytestmark = pytest.mark.asyncio
30
-
31
-
32
- @pytest.fixture()
33
- def available_nodes():
34
- nodes = {
35
- "0": IndexNode(
36
- id="0", address="node-0", available_disk=100, shard_count=0, dummy=True
37
- ),
38
- "1": IndexNode(
39
- id="1", address="node-1", available_disk=100, shard_count=0, dummy=True
40
- ),
41
- "2": IndexNode(
42
- id="2", address="node-2", available_disk=100, shard_count=0, dummy=True
43
- ),
44
- }
45
- with patch.object(manager, "INDEX_NODES", new=nodes):
46
- yield nodes
47
-
48
-
49
- @pytest.fixture()
50
- def shards():
51
- yield writer_pb2.Shards(
52
- shards=[
53
- writer_pb2.ShardObject(
54
- shard="1",
55
- replicas=[
56
- writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="1")),
57
- writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="2")),
58
- ],
59
- read_only=True,
60
- ),
61
- writer_pb2.ShardObject(
62
- shard="2",
63
- replicas=[
64
- writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="3")),
65
- writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="4")),
66
- ],
67
- read_only=False,
68
- ),
69
- ],
70
- kbid="kbid",
71
- actual=1,
72
- )
73
-
74
-
75
- @pytest.fixture()
76
- def resource_ids():
77
- yield ["1", "2", "3"]
78
-
79
-
80
- @pytest.fixture()
81
- def resources_datamanager(resource_ids):
82
- mock = MagicMock()
83
-
84
- async def iterate_resource_ids(kbid):
85
- for id in resource_ids:
86
- yield id
87
-
88
- mock.iterate_resource_ids = iterate_resource_ids
89
- mock.get_resource_shard_id = AsyncMock()
90
- mock.get_resource_shard_id.return_value = "1"
91
-
92
- res = MagicMock()
93
-
94
- res.basic.modified.ToDatetime.return_value = datetime.now()
95
-
96
- mock.get_resource = AsyncMock()
97
- mock.get_resource.return_value = res
98
-
99
- mock.get_resource_index_message = AsyncMock()
100
- metadata = MagicMock()
101
- metadata.modified.ToDatetime.return_value = datetime.now()
102
- mock.get_resource_index_message.return_value = metadata
103
-
104
- with patch("nucliadb.common.cluster.rollover.datamanagers.resources", mock):
105
- yield mock
106
-
107
-
108
- @pytest.fixture()
109
- def cluster_datamanager(resource_ids, shards):
110
- mock = MagicMock()
111
- mock.get_kb_shards = AsyncMock()
112
- mock.get_kb_shards.return_value = shards
113
- mock.update_kb_shards = AsyncMock()
114
-
115
- with patch("nucliadb.common.cluster.rollover.datamanagers.cluster", mock):
116
- yield mock
117
-
118
-
119
- @pytest.fixture()
120
- def rollover_datamanager(resource_ids, cluster_datamanager):
121
- mock = MagicMock()
122
- mock.get_kb_rollover_shards = AsyncMock()
123
- mock.get_kb_rollover_shards.return_value = None
124
- mock.update_kb_rollover_shards = AsyncMock()
125
- mock.delete_kb_rollover_shard = AsyncMock()
126
- mock.delete_kb_rollover_shards = AsyncMock()
127
- mock.get_to_index = AsyncMock(side_effect=["1", None])
128
- mock.add_indexed = AsyncMock()
129
- mock.remove_to_index = AsyncMock()
130
- mock.get_indexed_data = AsyncMock(return_value=("1", 1))
131
- mock.remove_indexed = AsyncMock()
132
-
133
- async def _mock_indexed_keys(kbid):
134
- yield "1"
135
-
136
- mock.iter_indexed_keys = _mock_indexed_keys
137
-
138
- with (
139
- patch("nucliadb.common.cluster.rollover.datamanagers.rollover", mock),
140
- patch(
141
- "nucliadb.common.cluster.rollover.datamanagers.with_transaction",
142
- return_value=AsyncMock(),
143
- ),
144
- patch(
145
- "nucliadb.ingest.consumer.shard_creator.locking.distributed_lock",
146
- return_value=AsyncMock(),
147
- ),
148
- ):
149
- yield mock
150
-
151
-
152
- @pytest.fixture()
153
- def app_context(rollover_datamanager, resources_datamanager, available_nodes):
154
- mock = MagicMock()
155
- mock.shard_manager = MagicMock()
156
- mock.shard_manager.rollback_shard = AsyncMock()
157
- mock.shard_manager.add_resource = AsyncMock()
158
- mock.shard_manager.delete_resource = AsyncMock()
159
- mock.kv_driver = MagicMock()
160
-
161
- consumer_info = MagicMock()
162
- consumer_info.delivered.stream_seq = 0
163
- mock.nats_manager.js.consumer_info.return_value = consumer_info
164
- yield mock
165
-
166
-
167
- async def test_create_rollover_shards(
168
- app_context, available_nodes, shards: writer_pb2.Shards, rollover_datamanager
169
- ):
170
- new_shards = await rollover.create_rollover_shards(app_context, "kbid")
171
-
172
- assert new_shards.kbid == "kbid"
173
- assert sum(
174
- [len(node.writer.calls["NewShard"]) for node in available_nodes.values()]
175
- ) == sum([len(s.replicas) for s in shards.shards])
176
- rollover_datamanager.update_kb_rollover_shards.assert_called_with(
177
- ANY, kbid="kbid", kb_shards=new_shards
178
- )
179
-
180
-
181
- async def test_create_rollover_shards_does_not_recreate(
182
- app_context, shards: writer_pb2.Shards, rollover_datamanager
183
- ):
184
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
185
-
186
- await rollover.create_rollover_shards(app_context, "kbid")
187
-
188
- app_context.shard_manager.rollback_shard.assert_not_called()
189
- rollover_datamanager.update_kb_rollover_shards.assert_not_called()
190
-
191
-
192
- async def test_index_rollover_shards(
193
- app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
194
- ):
195
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
196
-
197
- await rollover.index_rollover_shards(app_context, "kbid")
198
- rollover_datamanager.add_indexed.assert_called_with(
199
- ANY, kbid="kbid", resource_id="1", shard_id="1", modification_time=1
200
- )
201
-
202
-
203
- async def test_index_rollover_shards_handles_missing_shards(
204
- app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
205
- ):
206
- rollover_datamanager.get_kb_rollover_shards.return_value = None
207
- with pytest.raises(rollover.UnexpectedRolloverError):
208
- await rollover.index_rollover_shards(app_context, "kbid")
209
-
210
-
211
- async def test_index_rollover_shards_handles_missing_shard_id(
212
- app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
213
- ):
214
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
215
- resources_datamanager.get_resource_shard_id.return_value = None
216
- await rollover.index_rollover_shards(app_context, "kbid")
217
-
218
-
219
- async def test_index_rollover_shards_handles_missing_res(
220
- app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
221
- ):
222
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
223
- resources_datamanager.get_resource_index_message.return_value = None
224
-
225
- await rollover.index_rollover_shards(app_context, "kbid")
226
-
227
- rollover_datamanager.remove_to_index.assert_called_with(
228
- ANY, kbid="kbid", resource="1"
229
- )
230
-
231
-
232
- async def test_cutover_shards(
233
- app_context, rollover_datamanager, cluster_datamanager, shards
234
- ):
235
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
236
-
237
- await rollover.cutover_shards(app_context, "kbid")
238
-
239
- cluster_datamanager.update_kb_shards.assert_called_with(
240
- ANY, kbid="kbid", shards=ANY
241
- )
242
- [
243
- app_context.shard_manager.rollback_shard.assert_any_call(shard)
244
- for shard in shards.shards
245
- ]
246
-
247
-
248
- async def test_cutover_shards_missing(app_context, rollover_datamanager):
249
- rollover_datamanager.get_kb_rollover_shards.return_value = None
250
-
251
- with pytest.raises(rollover.UnexpectedRolloverError):
252
- await rollover.cutover_shards(app_context, "kbid")
253
-
254
-
255
- async def test_validate_indexed_data(
256
- app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
257
- ):
258
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
259
-
260
- indexed_res = await rollover.validate_indexed_data(app_context, "kbid")
261
- assert len(indexed_res) == len(resource_ids)
262
- [
263
- resources_datamanager.get_resource_index_message.assert_any_call(
264
- ANY, kbid="kbid", rid=res_id
265
- )
266
- for res_id in resource_ids
267
- ]
268
-
269
-
270
- async def test_validate_indexed_data_handles_missing_res(
271
- app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
272
- ):
273
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
274
- resources_datamanager.get_resource_index_message.return_value = None
275
- assert len(await rollover.validate_indexed_data(app_context, "kbid")) == 0
276
-
277
-
278
- async def test_rollover_shards(app_context, rollover_datamanager, shards, resource_ids):
279
- rollover_datamanager.get_kb_rollover_shards.return_value = shards
280
- resource_ids.clear()
281
-
282
- await rollover.rollover_kb_shards(app_context, "kbid")
@@ -1,18 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
@@ -1,127 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
19
- #
20
- import asyncio
21
- from unittest import mock
22
-
23
- import pytest
24
-
25
- from nucliadb.common.maindb.driver import Driver
26
-
27
-
28
- class TransactionTest:
29
- def __init__(self):
30
- self.abort = mock.AsyncMock()
31
- self.open = True
32
-
33
- async def commit(self, **kw):
34
- self.open = False
35
-
36
-
37
- @pytest.fixture(scope="function")
38
- def txn():
39
- return TransactionTest()
40
-
41
-
42
- @pytest.fixture(scope="function")
43
- def driver(txn) -> Driver: # type: ignore
44
- driver = Driver()
45
- with mock.patch.object(driver, "begin", new=mock.AsyncMock(return_value=txn)):
46
- yield driver
47
-
48
-
49
- @pytest.mark.asyncio
50
- async def test_driver_async_abort(driver, txn):
51
- async with driver.transaction(wait_for_abort=False):
52
- pass
53
-
54
- assert len(driver._abort_tasks) == 1
55
- await asyncio.sleep(0.1)
56
-
57
- txn.abort.assert_called_once()
58
- assert len(driver._abort_tasks) == 0
59
-
60
-
61
- @pytest.mark.asyncio
62
- async def test_driver_finalize_aborts_transactions(driver, txn):
63
- async with driver.transaction(wait_for_abort=False):
64
- pass
65
-
66
- assert len(driver._abort_tasks) == 1
67
-
68
- await driver.finalize()
69
-
70
- txn.abort.assert_called_once()
71
-
72
- assert len(driver._abort_tasks) == 0
73
-
74
-
75
- @pytest.mark.asyncio
76
- async def test_transaction_handles_txn_begin_errors(driver):
77
- driver.begin.side_effect = ValueError()
78
- testmock = mock.AsyncMock()
79
- with pytest.raises(ValueError):
80
- async with driver.transaction():
81
- testmock()
82
- testmock.assert_not_called()
83
-
84
-
85
- @pytest.mark.asyncio
86
- async def test_transaction_does_not_abotr_if_commited(driver):
87
- async with driver.transaction() as txn:
88
- await txn.commit()
89
- txn.abort.assert_not_awaited()
90
-
91
-
92
- @pytest.mark.asyncio
93
- async def test_transaction_aborts_if_txn_open(driver):
94
- async with driver.transaction() as txn:
95
- pass
96
- txn.abort.assert_awaited_once()
97
-
98
-
99
- @pytest.mark.asyncio
100
- async def test_transaction_aborts_on_errors(driver):
101
- with pytest.raises(ValueError):
102
- async with driver.transaction() as txn:
103
- raise ValueError()
104
- txn.abort.assert_awaited_once()
105
-
106
-
107
- @pytest.mark.asyncio
108
- async def test_transaction_wait_for_abort(driver):
109
- async with driver.transaction(wait_for_abort=False) as txn:
110
- pass
111
- txn.abort.assert_not_awaited()
112
- txn.abort.assert_called_once()
113
- txn.abort.reset_mock()
114
-
115
- # Should wait for abort if there are errors
116
- with pytest.raises(ValueError):
117
- async with driver.transaction(wait_for_abort=False) as txn:
118
- raise ValueError()
119
- txn.abort.assert_awaited_once()
120
- txn.abort.assert_called_once()
121
- txn.abort.reset_mock()
122
-
123
- async with driver.transaction(wait_for_abort=True) as txn:
124
- pass
125
- txn.abort.assert_awaited_once()
126
- txn.abort.assert_called_once()
127
- txn.abort.reset_mock()
@@ -1,53 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
19
-
20
- from unittest.mock import AsyncMock, MagicMock
21
-
22
- import pytest
23
-
24
- from nucliadb.common.maindb.exceptions import ConflictError
25
- from nucliadb.common.maindb.tikv import LeaderNotFoundError, TiKVDataLayer
26
-
27
-
28
- @pytest.mark.parametrize(
29
- "tikv_exception,handled_exception",
30
- [
31
- (
32
- Exception("gRPC error: RpcFailure: 4-DEADLINE_EXCEEDED Deadline Exceeded"),
33
- TimeoutError,
34
- ),
35
- (Exception("Leader of region 34234 is not found"), LeaderNotFoundError),
36
- ],
37
- )
38
- async def test_get_retrials(tikv_exception, handled_exception):
39
- inner_txn = MagicMock(get=AsyncMock(side_effect=tikv_exception))
40
- tikv_txn = TiKVDataLayer(inner_txn)
41
-
42
- with pytest.raises(handled_exception):
43
- await tikv_txn.get("key")
44
-
45
- assert inner_txn.get.call_count == 2
46
-
47
-
48
- async def test_commit_raises_conflict_error():
49
- inner_txn = MagicMock(commit=AsyncMock(side_effect=Exception("WriteConflict")))
50
- tikv_txn = TiKVDataLayer(inner_txn)
51
-
52
- with pytest.raises(ConflictError):
53
- await tikv_txn.commit()
@@ -1,92 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
19
-
20
- from unittest.mock import AsyncMock, patch
21
-
22
- import pytest
23
-
24
- from nucliadb.common.maindb.utils import settings, setup_driver
25
- from nucliadb.ingest.settings import DriverConfig
26
- from nucliadb_utils.exceptions import ConfigurationError
27
- from nucliadb_utils.utilities import Utility, clean_utility
28
-
29
-
30
- @pytest.fixture(autouse=True)
31
- def reset_driver_utils():
32
- clean_utility(Utility.MAINDB_DRIVER)
33
- yield
34
- clean_utility(Utility.MAINDB_DRIVER)
35
-
36
-
37
- @pytest.mark.asyncio
38
- async def test_setup_driver_redis():
39
- mock = AsyncMock(initialized=False)
40
- with (
41
- patch.object(settings, "driver", DriverConfig("redis")),
42
- patch.object(settings, "driver_redis_url", "driver_redis_url"),
43
- patch("nucliadb.common.maindb.utils.RedisDriver", return_value=mock),
44
- ):
45
- assert await setup_driver() == mock
46
- mock.initialize.assert_awaited_once()
47
-
48
-
49
- @pytest.mark.asyncio
50
- async def test_setup_driver_tikv():
51
- mock = AsyncMock(initialized=False)
52
- with (
53
- patch.object(settings, "driver", DriverConfig("tikv")),
54
- patch.object(settings, "driver_tikv_url", "driver_tikv_url"),
55
- patch("nucliadb.common.maindb.utils.TiKVDriver", return_value=mock),
56
- ):
57
- assert await setup_driver() == mock
58
- mock.initialize.assert_awaited_once()
59
-
60
-
61
- @pytest.mark.asyncio
62
- async def test_setup_driver_pg():
63
- mock = AsyncMock(initialized=False)
64
- with (
65
- patch.object(settings, "driver", DriverConfig("pg")),
66
- patch.object(settings, "driver_pg_url", "driver_pg_url"),
67
- patch("nucliadb.common.maindb.utils.PGDriver", return_value=mock),
68
- ):
69
- assert await setup_driver() == mock
70
- mock.initialize.assert_awaited_once()
71
-
72
-
73
- @pytest.mark.asyncio
74
- async def test_setup_driver_local():
75
- mock = AsyncMock(initialized=False)
76
- with (
77
- patch.object(settings, "driver", DriverConfig("local")),
78
- patch.object(settings, "driver_local_url", "driver_local_url"),
79
- patch("nucliadb.common.maindb.utils.LocalDriver", return_value=mock),
80
- ):
81
- assert await setup_driver() == mock
82
- mock.initialize.assert_awaited_once()
83
-
84
-
85
- @pytest.mark.asyncio
86
- async def test_setup_driver_error():
87
- with (
88
- patch.object(settings, "driver", DriverConfig("pg")),
89
- patch.object(settings, "driver_pg_url", None),
90
- pytest.raises(ConfigurationError),
91
- ):
92
- await setup_driver()
@@ -1,36 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
19
- #
20
- import asyncio
21
- from unittest.mock import AsyncMock
22
-
23
- from nucliadb.common.context import ApplicationContext
24
-
25
-
26
- async def test_initialize_happens_only_once():
27
- context = ApplicationContext()
28
- context._initialize = AsyncMock() # noqa
29
-
30
- tasks = []
31
- for _ in range(10):
32
- tasks.append(context.initialize())
33
- await asyncio.gather(*tasks)
34
-
35
- context._initialize.assert_awaited_once()
36
- context._initialized is True
@@ -1,19 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
19
- #
@@ -1,37 +0,0 @@
1
- # Copyright (C) 2021 Bosutech XXI S.L.
2
- #
3
- # nucliadb is offered under the AGPL v3.0 and as commercial software.
4
- # For commercial licensing, contact us at info@nuclia.com.
5
- #
6
- # AGPL:
7
- # This program is free software: you can redistribute it and/or modify
8
- # it under the terms of the GNU Affero General Public License as
9
- # published by the Free Software Foundation, either version 3 of the
10
- # License, or (at your option) any later version.
11
- #
12
- # This program is distributed in the hope that it will be useful,
13
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
14
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15
- # GNU Affero General Public License for more details.
16
- #
17
- # You should have received a copy of the GNU Affero General Public License
18
- # along with this program. If not, see <http://www.gnu.org/licenses/>.
19
- #
20
- from unittest.mock import AsyncMock, Mock
21
-
22
- from nucliadb.export_import.datamanager import ExportImportDataManager
23
-
24
-
25
- async def test_try_delete_from_storage():
26
- driver = Mock()
27
- storage = Mock()
28
-
29
- dm = ExportImportDataManager(driver, storage)
30
- dm.delete_export = AsyncMock(side_effect=ValueError()) # type: ignore
31
- dm.delete_import = AsyncMock(side_effect=ValueError()) # type: ignore
32
-
33
- await dm.try_delete_from_storage("export", "kbid", "foo")
34
- await dm.try_delete_from_storage("import", "kbid", "foo")
35
-
36
- dm.delete_export.assert_called()
37
- dm.delete_import.assert_called()