Save Dataset and DatasetElements in cache database

Closes #8 (closed)

tests/test_worker.py

 # -*- coding: utf-8 -*-
 
+import json
 from argparse import Namespace
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from arkindex_worker.cache import (
     CachedClassification,
+    CachedDataset,
+    CachedDatasetElement,
     CachedElement,
     CachedEntity,
     CachedImage,
@@ -30,15 +33,24 @@ def test_process_split(tmp_path, downloaded_images):
     first_page_id = UUID("e26e6803-18da-4768-be30-a0a68132107c")
     second_page_id = UUID("c673bd94-96b1-4a2e-8662-a4d806940b5f")
 
+    # The dataset should already be saved in database when we call `process_split`
+    cached_dataset = CachedDataset.create(
+        id=uuid4(),
+        name="My dataset",
+        state="complete",
+        sets=json.dumps(["train", "val", "test"]),
+    )
+
     worker.process_split(
         "train",
         [
             retrieve_element(first_page_id),
             retrieve_element(second_page_id),
         ],
+        cached_dataset,
     )
 
-    # Should have created 20 elements in total
+    # Should have created 19 elements in total
     assert CachedElement.select().count() == 19
 
     # Should have created two pages at root
@@ -125,6 +137,18 @@ def test_process_split(tmp_path, downloaded_images):
     assert tr_entity.confidence == 1.0
     assert tr_entity.worker_run_id is None
 
+    # Should have linked all the elements to the correct dataset & split
+    assert CachedDatasetElement.select().count() == 19
+    assert (
+        CachedDatasetElement.select()
+        .where(
+            CachedDatasetElement.dataset == cached_dataset,
+            CachedDatasetElement.set_name == "train",
+        )
+        .count()
+        == 19
+    )
+
     # Full structure of the archive
     assert sorted(tmp_path.rglob("*")) == [
         tmp_path / "db.sqlite",