From 241c527fb3ce57c5dc5821efd36996d24946558a Mon Sep 17 00:00:00 2001
From: John Wilkie <john@v7labs.com>
Date: Tue, 16 Jul 2024 14:47:19 +0100
Subject: [PATCH] Updated retry logic so that it worked with storing release
 statuses as enums

---
 darwin/cli_functions.py          | 2 +-
 darwin/dataset/release.py        | 2 +-
 darwin/dataset/remote_dataset.py | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/darwin/cli_functions.py b/darwin/cli_functions.py
index 199d6f4fe..395c7c829 100644
--- a/darwin/cli_functions.py
+++ b/darwin/cli_functions.py
@@ -468,7 +468,7 @@ def pull_dataset(
         print_new_version_info(client)
     except NotFound:
         _error(
-            f"Version '{dataset.identifier}:{version}' does not exist "
+            f"Version '{dataset.identifier}:{version}' does not exist. "
             f"Use 'darwin dataset releases' to list all available versions."
         )
     except UnsupportedExportFormat as uef:
diff --git a/darwin/dataset/release.py b/darwin/dataset/release.py
index 517ca91ef..54e32fbb7 100644
--- a/darwin/dataset/release.py
+++ b/darwin/dataset/release.py
@@ -94,7 +94,7 @@ def __init__(
         self.team_slug = team_slug
         self.version = version
         self.name = name
-        self.status = status.value
+        self.status = ReleaseStatus(status)
         self.url = url
         self.export_date = export_date
         self.image_count = image_count
diff --git a/darwin/dataset/remote_dataset.py b/darwin/dataset/remote_dataset.py
index 07bfe5879..c96ef194b 100644
--- a/darwin/dataset/remote_dataset.py
+++ b/darwin/dataset/remote_dataset.py
@@ -274,16 +274,16 @@ def pull(
         if release.format != "json" and release.format != "darwin_json_2":
             raise UnsupportedExportFormat(release.format)
 
-        if release.status == "pending":
+        if release.status.value == "pending":
             if retry:
-                while release.status == "pending" and retry_timeout > 0:
+                while release.status.value == "pending" and retry_timeout > 0:
                     console.print(
                         f"Release '{release.name}' for dataset '{self.name}' is still processing. Retrying in {retry_interval} seconds... {retry_timeout} seconds left before timeout."
                     )
                     time.sleep(retry_interval)
                     retry_timeout -= retry_interval
                     release = self.get_release(release.name, include_unavailable=retry)
-                if release.status == "pending":
+                if release.status.value == "pending":
                     raise ValueError(
                         f"Release {release.name} for dataset '{self.name}' is still processing. Please try again later."
                     )