From 7257f1768f3a9a600e0361d64625611d3a67c306 Mon Sep 17 00:00:00 2001
From: Arnaud <arnaud@status.im>
Date: Tue, 4 Feb 2025 19:29:07 +0100
Subject: [PATCH] Return bytes to availability when finished

---
 codex/sales/states/finished.nim           | 5 ++++-
 tests/codex/sales/states/testfinished.nim | 8 ++++++--
 tests/codex/sales/testsales.nim           | 6 +++---
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/codex/sales/states/finished.nim b/codex/sales/states/finished.nim
index 151300d0..a5742256 100644
--- a/codex/sales/states/finished.nim
+++ b/codex/sales/states/finished.nim
@@ -32,5 +32,8 @@ method run*(state: SaleFinished, machine: Machine): Future[?State] {.async.} =
   info "Slot finished and paid out",
     requestId = data.requestId, slotIndex = data.slotIndex
 
+  if onClear =? agent.context.onClear:
+    onClear(request, data.slotIndex)
+
   if onCleanUp =? agent.onCleanUp:
-    await onCleanUp(returnedCollateral = state.returnedCollateral)
+    await onCleanUp(returnedCollateral = state.returnedCollateral, returnBytes = true)
diff --git a/tests/codex/sales/states/testfinished.nim b/tests/codex/sales/states/testfinished.nim
index 4b353014..69150418 100644
--- a/tests/codex/sales/states/testfinished.nim
+++ b/tests/codex/sales/states/testfinished.nim
@@ -26,6 +26,7 @@ asyncchecksuite "sales state 'finished'":
   var returnBytesWas = bool.none
   var reprocessSlotWas = bool.none
   var returnedCollateralValue = UInt256.none
+  var saleCleared = bool.none
 
   setup:
     market = MockMarket.new()
@@ -39,6 +40,8 @@ asyncchecksuite "sales state 'finished'":
     let context = SalesContext(market: market, clock: clock)
     agent = newSalesAgent(context, request.id, slotIndex, request.some)
     agent.onCleanUp = onCleanUp
+    agent.context.onClear = some proc(request: StorageRequest, idx: UInt256) =
+      saleCleared = some true
     state = SaleFinished(returnedCollateral: some currentCollateral)
 
   test "switches to cancelled state when request expires":
@@ -49,8 +52,9 @@ asyncchecksuite "sales state 'finished'":
     let next = state.onFailed(request)
     check !next of SaleFailed
 
-  test "calls onCleanUp with returnBytes = false, reprocessSlot = true, and returnedCollateral = currentCollateral":
+  test "calls onCleanUp with returnBytes = true, reprocessSlot = true, and returnedCollateral = currentCollateral":
     discard await state.run(agent)
-    check eventually returnBytesWas == some false
+    check eventually returnBytesWas == some true
     check eventually reprocessSlotWas == some false
     check eventually returnedCollateralValue == some currentCollateral
+    check eventually saleCleared == some true
diff --git a/tests/codex/sales/testsales.nim b/tests/codex/sales/testsales.nim
index 724f0737..02c8addb 100644
--- a/tests/codex/sales/testsales.nim
+++ b/tests/codex/sales/testsales.nim
@@ -354,14 +354,14 @@ asyncchecksuite "Sales":
     check eventually getAvailability().freeSize ==
       availability.freeSize - request.ask.slotSize
 
-  test "non-downloaded bytes are returned to availability once finished":
+  test "bytes are returned to availability once finished":
     var slotIndex = 0.u256
     sales.onStore = proc(
         request: StorageRequest, slot: UInt256, onBatch: BatchProc
     ): Future[?!void] {.async.} =
       slotIndex = slot
       let blk = bt.Block.new(@[1.byte]).get
-      await onBatch(@[blk])
+      await onBatch(blk.repeat(request.ask.slotSize.truncate(int)))
 
     let sold = newFuture[void]()
     sales.onSale = proc(request: StorageRequest, slotIndex: UInt256) =
@@ -377,7 +377,7 @@ asyncchecksuite "Sales":
     market.slotState[request.slotId(slotIndex)] = SlotState.Finished
     clock.advance(request.ask.duration.truncate(int64))
 
-    check eventually getAvailability().freeSize == origSize - 1
+    check eventually getAvailability().freeSize == origSize
 
   test "ignores download when duration not long enough":
     availability.duration = request.ask.duration - 1