From 2c623d0fab3402a2fc6f24f9286a34dbc4a51f2b Mon Sep 17 00:00:00 2001
From: John Wilkie <john@v7labs.com>
Date: Fri, 28 Jun 2024 15:11:06 +0100
Subject: [PATCH] Unit tests

---
 tests/darwin/dataset/remote_dataset_test.py | 43 +++++++++++++++++
 tests/fixtures.py                           | 51 ++++++++++++++++++++-
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/tests/darwin/dataset/remote_dataset_test.py b/tests/darwin/dataset/remote_dataset_test.py
index 630d199cb..a4d88adc9 100644
--- a/tests/darwin/dataset/remote_dataset_test.py
+++ b/tests/darwin/dataset/remote_dataset_test.py
@@ -813,6 +813,27 @@ def fake_download_zip(self, path):
                 )
                 assert metadata_path.exists()
 
+    def test_pull_raises_value_error_when_retry_is_true_and_release_is_none(
+        self, remote_dataset
+    ):
+        with pytest.raises(ValueError):
+            remote_dataset.pull(release=None, retry=True)
+
+    @patch("time.sleep", return_value=None)
+    def test_num_retries(self, mock_sleep, remote_dataset, pending_release):
+        with patch.object(remote_dataset, "get_release", return_value=pending_release):
+            with pytest.raises(ValueError):
+                remote_dataset.pull(release=pending_release, retry=True)
+            assert mock_sleep.call_count == 30  # 300 seconds / 10 seconds interval
+
+    @patch("time.sleep", return_value=None)
+    def test_raises_after_max_retry_duration(
+        self, mock_sleep, remote_dataset, pending_release
+    ):
+        with patch.object(remote_dataset, "get_release", return_value=pending_release):
+            with pytest.raises(ValueError, match="is still processing after"):
+                remote_dataset.pull(release=pending_release, retry=True)
+
 
 class TestPullNamingConvention:
     def _test_pull_naming_convention(
@@ -1321,3 +1342,25 @@ def test_register_files_with_blocked_items(self, remote_dataset: RemoteDatasetV2
         )
         assert len(result["registered"]) == 0
         assert len(result["blocked"]) == 1
+
+
+@pytest.mark.usefixtures("file_read_write_test")
+class TestGetReleases:
+    @patch("darwin.backend_v2.BackendV2.get_exports")
+    def test_returns_unavailable_releases_when_retry_is_true(
+        self, mock_get_exports, remote_dataset, releases_api_response
+    ):
+        mock_get_exports.return_value = releases_api_response
+        releases = remote_dataset.get_releases(retry=True)
+        assert len(releases) == 2
+        assert isinstance(releases[0], Release)
+        assert isinstance(releases[1], Release)
+
+    @patch("darwin.backend_v2.BackendV2.get_exports")
+    def test_omits_unavailable_releases_when_retry_is_false(
+        self, mock_get_exports, remote_dataset, releases_api_response
+    ):
+        mock_get_exports.return_value = releases_api_response
+        releases = remote_dataset.get_releases(retry=False)
+        assert len(releases) == 1
+        assert isinstance(releases[0], Release)
diff --git a/tests/fixtures.py b/tests/fixtures.py
index 37059da55..87ee2029a 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -1,11 +1,13 @@
 import shutil
+from datetime import datetime
 from pathlib import Path
-from typing import Generator
+from typing import Generator, List
 from zipfile import ZipFile
 
 import pytest
 
 from darwin.config import Config
+from darwin.dataset.release import Release
 
 
 @pytest.fixture
@@ -118,3 +120,50 @@ def local_config_file(
     shutil.rmtree(darwin_path)
     if backup_darwin_path.exists():
         shutil.move(backup_darwin_path, darwin_path)
+
+
+@pytest.fixture
+def pending_release() -> Release:
+    return Release(
+        dataset_slug="dataset_slug",
+        team_slug="team_slug",
+        version="1",
+        name="name",
+        status="pending",
+        url=None,
+        export_date=datetime.now(),
+        image_count=1,
+        class_count=1,
+        available=False,
+        latest=False,
+        format="json",
+    )
+
+
+@pytest.fixture
+def releases_api_response() -> List[dict]:
+    return [
+        {
+            "name": "release_1",
+            "status": "complete",
+            "version": 1,
+            "format": "darwin_json_2",
+            "metadata": {
+                "num_images": 1,
+                "annotation_classes": [{"id": 1, "name": "test_class"}],
+            },
+            "inserted_at": "2024-06-28T12:29:02Z",
+            "latest": True,
+            "download_url": "download_url",
+        },
+        {
+            "name": "release_2",
+            "status": "pending",
+            "version": 2,
+            "format": "darwin_json_2",
+            "metadata": {},
+            "inserted_at": "2024-06-28T12:32:33Z",
+            "latest": False,
+            "download_url": None,
+        },
+    ]