Save Dataset and DatasetElements in cache database

Closes #8 (closed)

tests/test_worker.py

@@ -34,7 +34,7 @@ def test_process_split(tmp_path, downloaded_images):
     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(
+    worker.cached_dataset = CachedDataset.create(
         id=uuid4(),
         name="My dataset",
         state="complete",
@@ -47,7 +47,6 @@ def test_process_split(tmp_path, downloaded_images):
             retrieve_element(first_page_id),
             retrieve_element(second_page_id),
         ],
-        cached_dataset,
     )
 
     # Should have created 19 elements in total
@@ -142,7 +141,7 @@ def test_process_split(tmp_path, downloaded_images):
     assert (
         CachedDatasetElement.select()
         .where(
-            CachedDatasetElement.dataset == cached_dataset,
+            CachedDatasetElement.dataset == worker.cached_dataset,
             CachedDatasetElement.set_name == "train",
         )
         .count()