From f2ad9e1198c3e16276c47ff64d119ca474f381c7 Mon Sep 17 00:00:00 2001
From: Eric <5089238+emizzle@users.noreply.github.com>
Date: Thu, 14 Sep 2023 14:46:03 +1000
Subject: [PATCH] Add `WaitForStorageContractFinished` for marketplace tests

---
 DistTestCore/Marketplace/MarketplaceAccess.cs | 14 +++++++++++++-
 Utils/ByteSize.cs                             |  9 +++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/DistTestCore/Marketplace/MarketplaceAccess.cs b/DistTestCore/Marketplace/MarketplaceAccess.cs
index 0671582..5c05266 100644
--- a/DistTestCore/Marketplace/MarketplaceAccess.cs
+++ b/DistTestCore/Marketplace/MarketplaceAccess.cs
@@ -187,6 +187,18 @@ namespace DistTestCore.Marketplace
             WaitForStorageContractState(timeout, "finished");
         }
 
+        public void WaitForStorageContractFinished(ByteSize contractFileSize)
+        {
+            if (!contractStartUtc.HasValue)
+            {
+                WaitForStorageContractStarted(contractFileSize.ToTimeSpan());
+            }
+            var gracePeriod = TimeSpan.FromSeconds(10);
+            var currentContractTime = DateTime.UtcNow - contractStartUtc!.Value;
+            var timeout = (ContractDuration - currentContractTime) + gracePeriod;
+            WaitForStorageContractState(timeout, "finished");
+        }
+
         /// <summary>
         /// Wait for contract to start. Max timeout depends on contract filesize. Allows more time for larger files.
         /// </summary>
@@ -208,7 +220,7 @@ namespace DistTestCore.Marketplace
         {
             var lastState = "";
             var waitStart = DateTime.UtcNow;
-            
+
             log.Log($"Waiting for {Time.FormatDuration(timeout)} for contract '{PurchaseId}' to reach state '{desiredState}'.");
             while (lastState != desiredState)
             {
diff --git a/Utils/ByteSize.cs b/Utils/ByteSize.cs
index b069b04..170aaf7 100644
--- a/Utils/ByteSize.cs
+++ b/Utils/ByteSize.cs
@@ -10,6 +10,8 @@
 
         public long SizeInBytes { get; }
 
+        public const double DefaultSecondsPerMB = 10.0;
+
         public long ToMB()
         {
             return SizeInBytes / (1024 * 1024);
@@ -36,6 +38,12 @@
         {
             return Formatter.FormatByteSize(SizeInBytes);
         }
+
+        public TimeSpan ToTimeSpan(double secsPerMB = DefaultSecondsPerMB)
+        {
+            var filesizeInMb = SizeInBytes / (1024 * 1024);
+            return TimeSpan.FromSeconds(filesizeInMb * secsPerMB);
+        }
     }
 
     public static class ByteSizeIntExtensions
@@ -91,5 +99,6 @@
         {
             return Convert.ToInt64(i).TB();
         }
+
     }
 }