fix tests

Author: KirillKurdyukov

src/EFCore.Ydb/CHANGELOG.md

@@ -1,7 +1,6 @@
 - Added support for the YDB retry policy (ADO.NET) and new configuration methods in `YdbDbContextOptionsBuilder`:
   - `EnableRetryIdempotence()`: enables retries for errors classified as idempotent. You must ensure the operation itself is idempotent.
   - `UseRetryPolicy(YdbRetryPolicyConfig retryPolicyConfig)`: configures custom backoff parameters and the maximum number of retry attempts.
-  - `UseRetryPolicy(IRetryPolicy retryPolicy)`: advanced option giving full control over retry behavior. The provider does not manage the attempt counter; your implementation must handle it.
 
 ## v0.1.0
 

src/EFCore.Ydb/src/Infrastructure/YdbDbContextOptionsBuilder.cs

@@ -22,10 +22,7 @@ public YdbDbContextOptionsBuilder EnableRetryIdempotence()
         => UseRetryPolicy(new YdbRetryPolicyConfig { EnableRetryIdempotence = true });
 
     public YdbDbContextOptionsBuilder UseRetryPolicy(YdbRetryPolicyConfig retryPolicyConfig)
-        => ExecutionStrategy(d => new YdbExecutionStrategy(d, new YdbRetryPolicy(retryPolicyConfig)));
-
-    public YdbDbContextOptionsBuilder UseRetryPolicy(IRetryPolicy retryPolicy)
-        => ExecutionStrategy(d => new YdbExecutionStrategy(d, retryPolicy));
+        => ExecutionStrategy(d => new YdbExecutionStrategy(d, retryPolicyConfig));
 
     public YdbDbContextOptionsBuilder DisableRetryOnFailure() =>
         ExecutionStrategy(d => new NonRetryingExecutionStrategy(d));

src/EFCore.Ydb/src/Storage/Internal/YdbExecutionStrategy.cs

@@ -1,5 +1,6 @@
 using System;
 using Microsoft.EntityFrameworkCore.Storage;
+using Ydb.Sdk;
 using Ydb.Sdk.Ado;
 using Ydb.Sdk.Ado.RetryPolicy;
 
@@ -15,17 +16,23 @@ namespace EntityFrameworkCore.Ydb.Storage.Internal;
 /// so that exception types and ShouldRetryOn semantics match the provider.
 /// <br/>- This base <see cref="ExecutionStrategy"/> is a good place to emit metrics/logs (attempt number, delay, exception type, etc.).
 /// </summary>
-public class YdbExecutionStrategy(ExecutionStrategyDependencies dependencies, IRetryPolicy retryPolicy)
+public class YdbExecutionStrategy(ExecutionStrategyDependencies dependencies, YdbRetryPolicyConfig retryPolicyConfig)
 // We pass "placeholders" to the base class:
-// - int.MaxValue and TimeSpan.Zero are not used in the real retry logic.
+// - TimeSpan.Zero is not used in the real retry logic.
 // - Actual limits/delays are driven by IRetryPolicy.
-    : ExecutionStrategy(dependencies, int.MaxValue /* unused! */, TimeSpan.Zero /* unused! */)
+    : ExecutionStrategy(dependencies, retryPolicyConfig.MaxAttempts, TimeSpan.Zero /* unused! */)
 {
+    private readonly YdbRetryPolicy _retryPolicy = new(retryPolicyConfig);
+
     public override bool RetriesOnFailure => true;
 
-    protected override bool ShouldRetryOn(Exception exception) => exception is YdbException;
+    protected override bool ShouldRetryOn(Exception exception) =>
+        exception is YdbException ydbException &&
+        (ydbException.IsTransient || retryPolicyConfig.EnableRetryIdempotence && ydbException.Code is
+            StatusCode.ClientTransportUnknown or
+            StatusCode.ClientTransportUnavailable or
+            StatusCode.Undetermined);
 
-    protected override TimeSpan? GetNextDelay(Exception lastException) => lastException is YdbException ydbException
-        ? retryPolicy.GetNextDelay(ydbException, ExceptionsEncountered.Count - 1)
-        : null;
+    protected override TimeSpan? GetNextDelay(Exception lastException) =>
+        _retryPolicy.GetNextDelay((YdbException)lastException, ExceptionsEncountered.Count - 1);
 }

src/EFCore.Ydb/src/Storage/Internal/YdbExecutionStrategyFactory.cs

@@ -7,5 +7,5 @@ public class YdbExecutionStrategyFactory(ExecutionStrategyDependencies dependenc
     : RelationalExecutionStrategyFactory(dependencies)
 {
     protected override IExecutionStrategy CreateDefaultStrategy(ExecutionStrategyDependencies dependencies) =>
-        new YdbExecutionStrategy(dependencies, YdbRetryPolicy.Default);
+        new YdbExecutionStrategy(dependencies, YdbRetryPolicyConfig.Default);
 }

src/EFCore.Ydb/test/EntityFrameworkCore.Ydb.FunctionalTests/QueryExpressionInterceptionYdbTestBase.cs

@@ -29,7 +29,7 @@ public override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder build
         {
             new YdbDbContextOptionsBuilder(base.AddOptions(builder))
 #pragma warning disable EF1001
-                .ExecutionStrategy(d => new YdbExecutionStrategy(d, YdbRetryPolicy.Default));
+                .ExecutionStrategy(d => new YdbExecutionStrategy(d, YdbRetryPolicyConfig.Default));
 #pragma warning restore EF1001
             return builder;
         }

src/Ydb.Sdk/src/Ado/RetryPolicy/YdbRetryPolicy.cs

@@ -85,7 +85,7 @@ internal YdbRetryPolicy(YdbRetryPolicyConfig config, IRandom random) : this(conf
 
         return ydbException.Code switch
         {
-            StatusCode.BadSession or StatusCode.SessionBusy => TimeSpan.Zero,
+            StatusCode.BadSession or StatusCode.SessionBusy or StatusCode.SessionExpired => TimeSpan.Zero,
             StatusCode.Aborted or StatusCode.Undetermined =>
                 FullJitter(_fastBackoffBaseMs, _fastCapBackoffMs, _fastCeiling, attempt, _random),
             StatusCode.Unavailable or StatusCode.ClientTransportUnknown or StatusCode.ClientTransportUnavailable =>

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/RetryPolicy/YdbRetryPolicyTests.cs

@@ -10,6 +10,7 @@ public class YdbRetryPolicyTests
     [Theory]
     [InlineData(StatusCode.BadSession)]
     [InlineData(StatusCode.SessionBusy)]
+    [InlineData(StatusCode.SessionExpired)]
     public void GetNextDelay_WhenStatusIsBadSessionOrBusySession_ReturnTimeSpanZero(StatusCode statusCode)
     {
         var ydbRetryPolicy = new YdbRetryPolicy(new YdbRetryPolicyConfig { MaxAttempts = 2 });