@@ -5,6 +5,8 @@ import com.getcode.opencode.controllers.TransactionOperations
55import com.getcode.opencode.model.accounts.AccountCluster
66import com.getcode.opencode.model.accounts.AccountInfo
77import com.getcode.opencode.model.accounts.AccountType
8+ import com.getcode.opencode.model.financial.Fiat
9+ import com.getcode.solana.keys.Mint
810import com.getcode.solana.keys.PublicKey
911import io.mockk.coEvery
1012import io.mockk.coVerify
@@ -24,20 +26,28 @@ class UsdcDepositSweepTest {
2426 private val transactionOperations: TransactionOperations = mockk(relaxed = true )
2527 private val accountController: AccountController = mockk(relaxed = true )
2628 private val tokenCoordinator: TokenCoordinator = mockk(relaxed = true )
29+ private val balancePoller: BalancePoller = mockk(relaxed = true )
2730
2831 private val owner: AccountCluster = mockk(relaxed = true )
2932
3033 private lateinit var sweep: UsdcDepositSweep
3134
3235 @Before
3336 fun setUp () {
37+ coEvery {
38+ balancePoller.awaitBalanceChange(any(), any(), any(), any(), any())
39+ } returns Result .success(Fiat .Zero )
40+
3441 sweep = UsdcDepositSweep (
3542 transactionOperations = transactionOperations,
3643 accountController = accountController,
3744 tokenCoordinator = tokenCoordinator,
45+ balancePoller = balancePoller,
3846 maxRetries = 3 ,
3947 initialDelay = 10 .milliseconds,
4048 backoffFactor = 1.0 ,
49+ pollInterval = 10 .milliseconds,
50+ pollMaxAttempts = 2 ,
4151 )
4252 }
4353
@@ -144,7 +154,7 @@ class UsdcDepositSweepTest {
144154 }
145155
146156 @Test
147- fun `calls tokenCoordinator update on successful swap` () = runTest {
157+ fun `calls tokenCoordinator update and polls for USDF balance on successful swap` () = runTest {
148158 stubUsdcAccount(balance = 1_000_000L )
149159 coEvery { transactionOperations.swapUsdc(any(), any()) } returns Result .success(Unit )
150160
@@ -155,6 +165,15 @@ class UsdcDepositSweepTest {
155165 coVerify {
156166 tokenCoordinator.update()
157167 }
168+ coVerify {
169+ balancePoller.awaitBalanceChange(
170+ mint = Mint .usdf,
171+ baseline = Fiat .Zero ,
172+ predicate = any(),
173+ interval = 10 .milliseconds,
174+ maxAttempts = 2 ,
175+ )
176+ }
158177 }
159178
160179 @Test
@@ -173,6 +192,42 @@ class UsdcDepositSweepTest {
173192 }
174193 }
175194
195+ @Test
196+ fun `does not poll for USDF balance when swap fails` () = runTest {
197+ stubUsdcAccount(balance = 1_000_000L )
198+ coEvery {
199+ transactionOperations.swapUsdc(any(), any())
200+ } returns Result .failure(RuntimeException (" swap failed" ))
201+
202+ sweep.execute(owner)
203+ advanceUntilIdle()
204+ Thread .sleep(200 )
205+
206+ coVerify(exactly = 0 ) {
207+ balancePoller.awaitBalanceChange(any(), any(), any(), any(), any())
208+ }
209+ }
210+
211+ @Test
212+ fun `completes gracefully when USDF balance poll times out` () = runTest {
213+ stubUsdcAccount(balance = 1_000_000L )
214+ coEvery { transactionOperations.swapUsdc(any(), any()) } returns Result .success(Unit )
215+ coEvery {
216+ balancePoller.awaitBalanceChange(any(), any(), any(), any(), any())
217+ } returns Result .failure(BalancePollError .Timeout (Mint .usdf, 2 ))
218+
219+ sweep.execute(owner)
220+ advanceUntilIdle()
221+ Thread .sleep(200 )
222+
223+ coVerify {
224+ tokenCoordinator.update()
225+ }
226+ coVerify {
227+ balancePoller.awaitBalanceChange(any(), any(), any(), any(), any())
228+ }
229+ }
230+
176231 @Test
177232 fun `does not execute concurrently when job is active` () = runTest {
178233 stubUsdcAccount(balance = 1_000_000L )
0 commit comments