Add ifdef to enable unittest on ObservableDbOperation and add unitests for it

src/Microsoft.AspNet.SignalR.SqlServer/DbOperation.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Data;
 using System.Data.Common;
 using System.Data.SqlClient;
 using System.Diagnostics.CodeAnalysis;
@@ -21,7 +22,7 @@ namespace Microsoft.AspNet.SignalR.SqlServer
         public DbOperation(string connectionString, string commandText, ILogger logger)
             : this(connectionString, commandText, logger, SqlClientFactory.Instance.AsIDbProviderFactory())
         {
-            
+
         }
 
         public DbOperation(string connectionString, string commandText, ILogger logger, IDbProviderFactory dbProviderFactory)
@@ -71,12 +72,20 @@ namespace Microsoft.AspNet.SignalR.SqlServer
             return tcs.Task;
         }
 
+#if ASPNET50
+        public virtual int ExecuteReader(Action<IDataRecord, DbOperation> processRecord)
+#else
         public virtual int ExecuteReader(Action<DbDataReader, DbOperation> processRecord)
+#endif
         {
             return ExecuteReader(processRecord, null);
         }
 
+#if ASPNET50
+        protected virtual int ExecuteReader(Action<IDataRecord, DbOperation> processRecord, Action<IDbCommand> commandAction)
+#else
         protected virtual int ExecuteReader(Action<DbDataReader, DbOperation> processRecord, Action<DbCommand> commandAction)
+#endif      
         {
             return Execute(cmd =>
             {
@@ -100,7 +109,11 @@ namespace Microsoft.AspNet.SignalR.SqlServer
 
         [SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope", Justification = "It's the caller's responsibility to dispose as the command is returned"),
          SuppressMessage("Microsoft.Security", "CA2100:Review SQL queries for security vulnerabilities", Justification = "General purpose SQL utility command")]
-        protected virtual DbCommand CreateCommand(DbConnection connection)
+#if ASPNET50
+        protected virtual IDbCommand CreateCommand(IDbConnection connection)
+#else
+        protected virtual DbCommand CreateCommand(DbConnection connection)                  
+#endif
         {
             var command = connection.CreateCommand();
             command.CommandText = CommandText;
@@ -117,11 +130,18 @@ namespace Microsoft.AspNet.SignalR.SqlServer
         }
 
         [SuppressMessage("Microsoft.Usage", "CA2202:Do not dispose objects multiple times", Justification = "False positive?")]
+#if ASPNET50
+        private T Execute<T>(Func<IDbCommand, T> commandFunc)
+#else
         private T Execute<T>(Func<DbCommand, T> commandFunc)
+#endif
         {
             T result = default(T);
+#if ASPNET50
+            IDbConnection connection = null;
+#else
             DbConnection connection = null;
-            
+#endif
             try
             {
                 connection = _dbProviderFactory.CreateConnection();
@@ -142,7 +162,11 @@ namespace Microsoft.AspNet.SignalR.SqlServer
             return result;
         }
 
+#if ASPNET50
+        private void LoggerCommand(IDbCommand command)
+#else
         private void LoggerCommand(DbCommand command)
+#endif
         {
             if (Logger.IsEnabled(LogLevel.Verbose))
             {
@@ -155,10 +179,17 @@ namespace Microsoft.AspNet.SignalR.SqlServer
 
         [SuppressMessage("Microsoft.Usage", "CA2202:Do not dispose objects multiple times", Justification = "Disposed in async Finally block"),
          SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope", Justification = "Disposed in async Finally block")]
+#if ASPNET50
+        private void Execute<T>(Func<IDbCommand, Task<T>> commandFunc, TaskCompletionSource<T> tcs)
+#else
         private void Execute<T>(Func<DbCommand, Task<T>> commandFunc, TaskCompletionSource<T> tcs)
+#endif
         {
+#if ASPNET50
+            IDbConnection connection = null;
+#else
             DbConnection connection = null;
-           
+#endif
             try
             {
                 connection = _dbProviderFactory.CreateConnection();

src/Microsoft.AspNet.SignalR.SqlServer/DbProviderFactoryAdapter.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Open Technologies, Inc. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System.Data;
 using System.Data.Common;
 
 namespace Microsoft.AspNet.SignalR.SqlServer
@@ -14,7 +15,11 @@ namespace Microsoft.AspNet.SignalR.SqlServer
             _dbProviderFactory = dbProviderFactory;
         }
 
+#if ASPNET50
+        public IDbConnection CreateConnection()
+#else
         public DbConnection CreateConnection()
+#endif
         {
             return _dbProviderFactory.CreateConnection();
         }

src/Microsoft.AspNet.SignalR.SqlServer/IDataRecordExtensions.cs

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Open Technologies, Inc. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Data;
+using System.Data.SqlClient;
+
+namespace Microsoft.AspNet.SignalR.SqlServer
+{
+    internal static class IDataRecordExtensions
+    {
+        public static byte[] GetBinary(this IDataRecord reader, int ordinalIndex)
+        {
+            var sqlReader = reader as SqlDataReader;
+            if (sqlReader == null)
+            {
+                throw new NotSupportedException();
+            }
+
+            return sqlReader.GetSqlBinary(ordinalIndex).Value;
+        }
+    }
+}

src/Microsoft.AspNet.SignalR.SqlServer/IDbBehavior.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Data;
 using System.Data.Common;
 using System.Data.SqlClient;
 
@@ -14,7 +15,7 @@ namespace Microsoft.AspNet.SignalR.SqlServer
         IList<Tuple<int, int>> UpdateLoopRetryDelays { get; }
 
 #if ASPNET50
-        void AddSqlDependency(DbCommand command, Action<SqlNotificationEventArgs> callback);
+        void AddSqlDependency(IDbCommand command, Action<SqlNotificationEventArgs> callback);
 #endif
     }
 }

src/Microsoft.AspNet.SignalR.SqlServer/IDbCommandExtensions.cs

@@ -2,11 +2,13 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Data;
 using System.Data.Common;
 using System.Data.SqlClient;
 using System.Threading.Tasks;
 using JetBrains.Annotations;
 
+
 namespace Microsoft.AspNet.SignalR.SqlServer
 {
     internal static class IDbCommandExtensions
@@ -14,7 +16,7 @@ namespace Microsoft.AspNet.SignalR.SqlServer
         private readonly static TimeSpan _dependencyTimeout = TimeSpan.FromSeconds(60);
 
 #if ASPNET50
-        public static void AddSqlDependency([NotNull]this DbCommand command, Action<SqlNotificationEventArgs> callback)
+        public static void AddSqlDependency([NotNull]this IDbCommand command, Action<SqlNotificationEventArgs> callback)
         {
             var sqlCommand = command as SqlCommand;
             if (sqlCommand == null)
@@ -26,7 +28,12 @@ namespace Microsoft.AspNet.SignalR.SqlServer
             dependency.OnChange += (o, e) => callback(e);
         }
 #endif
+
+#if ASPNET50
+        public static Task<int> ExecuteNonQueryAsync(this IDbCommand command)
+#else
         public static Task<int> ExecuteNonQueryAsync(this DbCommand command)
+#endif
         {
             var sqlCommand = command as SqlCommand;
 

src/Microsoft.AspNet.SignalR.SqlServer/IDbProviderFactory.cs

@@ -1,13 +1,18 @@
 // Copyright (c) Microsoft Open Technologies, Inc. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System.Data;
 using System.Data.Common;
 
 namespace Microsoft.AspNet.SignalR.SqlServer
 {
     public interface IDbProviderFactory
     {
+#if ASPNET50
+        IDbConnection CreateConnection();
+#else
         DbConnection CreateConnection();
+#endif
         DbParameter CreateParameter();
     }
 }

src/Microsoft.AspNet.SignalR.SqlServer/ObservableDbOperation.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Data;
 using System.Data.Common;
 using System.Data.SqlClient;
 using System.Diagnostics;
@@ -84,7 +85,11 @@ namespace Microsoft.AspNet.SignalR.SqlServer
         /// </summary>
         [SuppressMessage("Microsoft.Maintainability", "CA1502:AvoidExcessiveComplexity", Justification = "Needs refactoring"),
          SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "Errors are reported via the callback")]
+#if ASPNET50
+        public void ExecuteReaderWithUpdates(Action<IDataRecord, DbOperation> processRecord)
+#else
         public void ExecuteReaderWithUpdates(Action<DbDataReader, DbOperation> processRecord)
+#endif
         {
             lock (_stopLocker)
             {
@@ -271,7 +276,7 @@ namespace Microsoft.AspNet.SignalR.SqlServer
         }
 
 #if ASPNET50
-        protected virtual void AddSqlDependency(DbCommand command, Action<SqlNotificationEventArgs> callback)
+        protected virtual void AddSqlDependency(IDbCommand command, Action<SqlNotificationEventArgs> callback)
         {
             command.AddSqlDependency(e => callback(e));
         }
@@ -458,7 +463,7 @@ namespace Microsoft.AspNet.SignalR.SqlServer
         }
 
 #if ASPNET50
-        void IDbBehavior.AddSqlDependency(DbCommand command, Action<SqlNotificationEventArgs> callback)
+        void IDbBehavior.AddSqlDependency(IDbCommand command, Action<SqlNotificationEventArgs> callback)
         {
             AddSqlDependency(command, callback);
         }

src/Microsoft.AspNet.SignalR.SqlServer/SqlPayload.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Data;
 using System.Data.Common;
 using JetBrains.Annotations;
 using Microsoft.AspNet.SignalR.Messaging;
@@ -17,7 +18,11 @@ namespace Microsoft.AspNet.SignalR.SqlServer
             return message.ToBytes();
         }
 
+#if ASPNET50
+        public static ScaleoutMessage FromBytes(IDataRecord record)
+#else
         public static ScaleoutMessage FromBytes(DbDataReader record)
+#endif
         {
             var message = ScaleoutMessage.FromBytes(record.GetBinary(1));
 

src/Microsoft.AspNet.SignalR.SqlServer/SqlReceiver.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Data;
 using System.Data.Common;
 using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
@@ -88,14 +89,14 @@ namespace Microsoft.AspNet.SignalR.SqlServer
                     _lastPayloadId = (long?)lastPayloadIdOperation.ExecuteScalar();
                     Queried();
 
-                     _logger.WriteVerbose(String.Format("{0}SqlReceiver started, initial payload id={1}", _loggerPrefix, _lastPayloadId));
+                    _logger.WriteVerbose(String.Format("{0}SqlReceiver started, initial payload id={1}", _loggerPrefix, _lastPayloadId));
 
                     // Complete the StartReceiving task as we've successfully initialized the payload ID
                     tcs.TrySetResult(null);
                 }
                 catch (Exception ex)
                 {
-                     _logger.WriteError(String.Format("{0}SqlReceiver error starting: {1}", _loggerPrefix, ex));
+                    _logger.WriteError(String.Format("{0}SqlReceiver error starting: {1}", _loggerPrefix, ex));
 
                     tcs.TrySetException(ex);
                     return;
@@ -137,12 +138,16 @@ namespace Microsoft.AspNet.SignalR.SqlServer
             _logger.WriteInformation("{0}SqlReceiver.Receive returned", _loggerPrefix);
         }
 
+#if ASPNET50
+        private void ProcessRecord(IDataRecord record, DbOperation dbOperation)
+#else
         private void ProcessRecord(DbDataReader record, DbOperation dbOperation)
+#endif
         {
             var id = record.GetInt64(0);
             ScaleoutMessage message = SqlPayload.FromBytes(record);
 
-             _logger.WriteVerbose(String.Format("{0}SqlReceiver last payload ID={1}, new payload ID={2}", _loggerPrefix, _lastPayloadId, id));
+            _logger.WriteVerbose(String.Format("{0}SqlReceiver last payload ID={1}, new payload ID={2}", _loggerPrefix, _lastPayloadId, id));
 
             if (id > _lastPayloadId + 1)
             {
@@ -158,9 +163,9 @@ namespace Microsoft.AspNet.SignalR.SqlServer
             // Update the Parameter with the new payload ID
             dbOperation.Parameters[0].Value = _lastPayloadId;
 
-             _logger.WriteVerbose(String.Format("{0}Updated receive reader initial payload ID parameter={1}", _loggerPrefix, _dbOperation.Parameters[0].Value));
+            _logger.WriteVerbose(String.Format("{0}Updated receive reader initial payload ID parameter={1}", _loggerPrefix, _dbOperation.Parameters[0].Value));
 
-             _logger.WriteVerbose(String.Format("{0}Payload {1} containing {2} message(s) received", _loggerPrefix, id, message.Messages.Count));
+            _logger.WriteVerbose(String.Format("{0}Payload {1} containing {2} message(s) received", _loggerPrefix, id, message.Messages.Count));
 
             Received((ulong)id, message);
         }

src/Microsoft.AspNet.SignalR.SqlServer/project.json

@@ -2,15 +2,9 @@
     "version": "3.0.0-*",
     "description": "Core server components for ASP.NET SignalR.",
     "dependencies": {
-        "System.Data.Common": "1.0.0.0-*",
-        "System.Data.SqlClient": "1.0.0.0-*",
-        "Microsoft.AspNet.HttpFeature": { "version": "1.0.0-*", "type": "build" },
-        "Microsoft.AspNet.RequestContainer": "1.0.0-*",
-        "Microsoft.AspNet.Security.DataProtection": "1.0.0-*",
-        "Microsoft.AspNet.SignalR.Server": "3.0.0-*",
-        "Microsoft.Framework.Logging": "1.0.0-*",
-        "Microsoft.Framework.Runtime.Interfaces": { "version": "1.0.0-*", "type": "build" },
-        "Newtonsoft.Json": "6.0.6"
+        "System.Data.Common": "1.0.0-*",
+        "System.Data.SqlClient": "1.0.0-*",
+        "Microsoft.AspNet.SignalR.Server": "3.0.0-*"
     },
     "compilationOptions": {
         "warningsAsErrors": true

test/Microsoft.AspNet.SignalR.SqlServer.Tests/ObservableSqlOperationFacts.cs

@@ -0,0 +1,285 @@
+// Copyright (c) Microsoft Open Technologies, Inc. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Data;
+using System.Data.Common;
+using System.Data.SqlClient;
+using System.Linq;
+using System.Threading;
+using Microsoft.AspNet.SignalR.SqlServer;
+using Microsoft.Framework.Logging;
+using Moq;
+using Xunit;
+
+namespace Microsoft.AspNet.SignalR.Tests.SqlServer
+{
+    public class ObservableSqlOperationFacts
+    {
+        private static readonly List<Tuple<int, int>> _defaultRetryDelays = new List<Tuple<int, int>> { new Tuple<int, int>(0, 1) };
+
+        [Theory]
+        [InlineData(true)]
+        [InlineData(false)]
+        public void UseSqlNotificationsIfAvailable(bool supportSqlNotifications)
+        {
+            // Arrange
+            var sqlDependencyAdded = false;
+            var retryLoopCount = 0;
+            var mre = new ManualResetEventSlim();
+            var dbProviderFactory = new MockDbProviderFactory();
+            var dbBehavior = new Mock<IDbBehavior>();
+            var logger = new Mock<ILogger>();
+            dbBehavior.Setup(db => db.UpdateLoopRetryDelays).Returns(_defaultRetryDelays);
+            dbBehavior.Setup(db => db.StartSqlDependencyListener()).Returns(supportSqlNotifications);
+            dbBehavior.Setup(db => db.AddSqlDependency(It.IsAny<IDbCommand>(), It.IsAny<Action<SqlNotificationEventArgs>>()))
+                .Callback(() =>
+                {
+                    sqlDependencyAdded = true;
+                    mre.Set();
+                });
+            var operation = new ObservableDbOperation("test", "test", logger.Object, dbProviderFactory, dbBehavior.Object);
+            operation.Faulted += _ => mre.Set();
+            operation.Queried += () =>
+            {
+                retryLoopCount++;
+                if (retryLoopCount > 1)
+                {
+                    mre.Set();
+                }
+            };
+
+            // Act
+            ThreadPool.QueueUserWorkItem(_ => operation.ExecuteReaderWithUpdates((record, o) => { }));
+            mre.Wait();
+            operation.Dispose();
+
+            // Assert
+            Assert.Equal(supportSqlNotifications, sqlDependencyAdded);
+        }
+
+        [Theory]
+        [InlineData(1, null, null)]
+        [InlineData(5, null, null)]
+        [InlineData(10, null, null)]
+        [InlineData(1, 5, 10)]
+        public void DoesRetryLoopConfiguredNumberOfTimes(int? length1, int? length2, int? length3)
+        {
+            // Arrange
+            var retryLoopCount = 0;
+            var mre = new ManualResetEventSlim();
+            var retryLoopArgs = new List<int?>(new[] { length1, length2, length3 }).Where(l => l.HasValue);
+            var retryLoopTotal = retryLoopArgs.Sum().Value;
+            var retryLoopDelays = new List<Tuple<int, int>>(retryLoopArgs.Select(l => new Tuple<int, int>(0, l.Value)));
+            var sqlDependencyCreated = false;
+            var dbProviderFactory = new MockDbProviderFactory();
+            var dbBehavior = new Mock<IDbBehavior>();
+            var logger = new Mock<ILogger>();
+            dbBehavior.Setup(db => db.UpdateLoopRetryDelays).Returns(retryLoopDelays);
+            dbBehavior.Setup(db => db.StartSqlDependencyListener()).Returns(true);
+            dbBehavior.Setup(db => db.AddSqlDependency(It.IsAny<IDbCommand>(), It.IsAny<Action<SqlNotificationEventArgs>>()))
+                .Callback(() =>
+                {
+                    sqlDependencyCreated = true;
+                    mre.Set();
+                });
+            var operation = new ObservableDbOperation("test", "test", logger.Object, dbProviderFactory, dbBehavior.Object);
+            operation.Faulted += _ => mre.Set();
+            operation.Queried += () =>
+            {
+                if (!sqlDependencyCreated)
+                {
+                    // Only update the loop count if the SQL dependency hasn't been created yet (we're still in the loop)
+                    retryLoopCount++;
+                }
+                if (retryLoopCount == retryLoopTotal)
+                {
+                    mre.Set();
+                }
+            };
+
+            // Act
+            ThreadPool.QueueUserWorkItem(_ => operation.ExecuteReaderWithUpdates((record, o) => { }));
+            mre.Wait();
+            operation.Dispose();
+
+            // Assert
+            Assert.Equal(retryLoopTotal, retryLoopCount);
+        }
+
+        [Fact]
+        public void CallsOnErrorOnException()
+        {
+            // Arrange
+            var mre = new ManualResetEventSlim(false);
+            var onErrorCalled = false;
+            var dbProviderFactory = new MockDbProviderFactory();
+            var dbBehavior = new Mock<IDbBehavior>();
+            var logger = new Mock<ILogger>();
+            dbBehavior.Setup(db => db.UpdateLoopRetryDelays).Returns(_defaultRetryDelays);
+            dbBehavior.Setup(db => db.StartSqlDependencyListener()).Returns(false);
+            dbProviderFactory.MockDataReader.Setup(r => r.Read()).Throws(new ApplicationException("test"));
+            var operation = new ObservableDbOperation("test", "test", logger.Object, dbProviderFactory, dbBehavior.Object);
+            operation.Faulted += _ =>
+            {
+                onErrorCalled = true;
+                mre.Set();
+            };
+
+            // Act
+            ThreadPool.QueueUserWorkItem(_ => operation.ExecuteReaderWithUpdates((record, o) => { }));
+            mre.Wait();
+            operation.Dispose();
+
+            // Assert
+            Assert.True(onErrorCalled);
+        }
+
+        [Fact]
+        public void ExecuteReaderSetsNotificationStateCorrectlyUpToAwaitingNotification()
+        {
+            // Arrange
+            var retryLoopDelays = new[] { Tuple.Create(0, 1) };
+            var dbProviderFactory = new MockDbProviderFactory();
+            var dbBehavior = new Mock<IDbBehavior>();
+            var logger = new Mock<ILogger>();
+            dbBehavior.Setup(db => db.UpdateLoopRetryDelays).Returns(retryLoopDelays);
+            dbBehavior.Setup(db => db.StartSqlDependencyListener()).Returns(true);
+            dbBehavior.Setup(db => db.AddSqlDependency(It.IsAny<IDbCommand>(), It.IsAny<Action<SqlNotificationEventArgs>>()));
+            var operation = new ObservableDbOperation("test", "test", logger.Object, dbProviderFactory, dbBehavior.Object);
+            operation.Queried += () =>
+            {
+                // Currently in the query loop
+                Assert.Equal(ObservableDbOperation.NotificationState.ProcessingUpdates, operation.CurrentNotificationState);
+            };
+
+            // Act
+            operation.ExecuteReaderWithUpdates((_, __) => { });
+
+            // Assert
+            Assert.Equal(ObservableDbOperation.NotificationState.AwaitingNotification, operation.CurrentNotificationState);
+
+            operation.Dispose();
+        }
+
+        [Fact]
+        public void ExecuteReaderSetsNotificationStateCorrectlyWhenRecordsReceivedWhileSettingUpSqlDependency()
+        {
+            // Arrange
+            var mre = new ManualResetEventSlim(false);
+            var retryLoopDelays = new[] { Tuple.Create(0, 1) };
+            var dbProviderFactory = new MockDbProviderFactory();
+            var readCount = 0;
+            var sqlDependencyAddedCount = 0;
+            dbProviderFactory.MockDataReader.Setup(r => r.Read()).Returns(() => ++readCount == 2 && sqlDependencyAddedCount == 1);
+            var dbBehavior = new Mock<IDbBehavior>();
+            var logger = new Mock<ILogger>();
+            dbBehavior.Setup(db => db.UpdateLoopRetryDelays).Returns(retryLoopDelays);
+            dbBehavior.Setup(db => db.StartSqlDependencyListener()).Returns(true);
+            dbBehavior.Setup(db => db.AddSqlDependency(It.IsAny<IDbCommand>(), It.IsAny<Action<SqlNotificationEventArgs>>())).Callback(() => sqlDependencyAddedCount++);
+            var operation = new ObservableDbOperation("test", "test", logger.Object, dbProviderFactory, dbBehavior.Object);
+            long? stateOnLoopRestart = null;
+            var queriedCount = 0;
+            operation.Queried += () =>
+            {
+                queriedCount++;
+
+                if (queriedCount == 3)
+                {
+                    // First query after the loop starts again, check the state is reset
+                    stateOnLoopRestart = operation.CurrentNotificationState;
+                    mre.Set();
+                }
+            };
+
+            // Act
+            ThreadPool.QueueUserWorkItem(_ => operation.ExecuteReaderWithUpdates((__, ___) => { }));
+
+            mre.Wait();
+
+            Assert.True(stateOnLoopRestart.HasValue);
+            Assert.Equal(ObservableDbOperation.NotificationState.ProcessingUpdates, stateOnLoopRestart.Value);
+
+            operation.Dispose();
+        }
+
+        [Fact]
+        public void ExecuteReaderSetsNotificationStateCorrectlyWhenNotificationReceivedBeforeChangingStateToAwaitingNotification()
+        {
+            // Arrange
+            var mre = new ManualResetEventSlim(false);
+            var retryLoopDelays = new[] { Tuple.Create(0, 1) };
+            var dbProviderFactory = new MockDbProviderFactory();
+            var sqlDependencyAdded = false;
+            var dbBehavior = new Mock<IDbBehavior>();
+            var logger = new Mock<ILogger>();
+            dbBehavior.Setup(db => db.UpdateLoopRetryDelays).Returns(retryLoopDelays);
+            dbBehavior.Setup(db => db.StartSqlDependencyListener()).Returns(true);
+            dbBehavior.Setup(db => db.AddSqlDependency(It.IsAny<IDbCommand>(), It.IsAny<Action<SqlNotificationEventArgs>>()))
+                      .Callback(() => sqlDependencyAdded = true);
+            var operation = new ObservableDbOperation("test", "test", logger.Object, dbProviderFactory, dbBehavior.Object);
+            dbProviderFactory.MockDataReader.Setup(r => r.Read()).Returns(() =>
+            {
+                if (sqlDependencyAdded)
+                {
+                    // Fake the SQL dependency firing while we're setting it up
+                    operation.CurrentNotificationState = ObservableDbOperation.NotificationState.NotificationReceived;
+                    sqlDependencyAdded = false;
+                }
+                return false;
+            });
+            long? stateOnLoopRestart = null;
+            var queriedCount = 0;
+            operation.Queried += () =>
+            {
+                queriedCount++;
+
+                if (queriedCount == 3)
+                {
+                    // First query after the loop starts again, capture the state
+                    stateOnLoopRestart = operation.CurrentNotificationState;
+                    mre.Set();
+                }
+            };
+
+            // Act
+            ThreadPool.QueueUserWorkItem(_ => operation.ExecuteReaderWithUpdates((__, ___) => { }));
+
+            mre.Wait();
+
+            Assert.True(stateOnLoopRestart.HasValue);
+            Assert.Equal(ObservableDbOperation.NotificationState.ProcessingUpdates, stateOnLoopRestart.Value);
+
+            operation.Dispose();
+        }
+
+        private class MockDbProviderFactory : IDbProviderFactory
+        {
+            public MockDbProviderFactory()
+            {
+                MockDbConnection = new Mock<IDbConnection>();
+                MockDbCommand = new Mock<IDbCommand>();
+                MockDataReader = new Mock<IDataReader>();
+
+                MockDbConnection.Setup(c => c.CreateCommand()).Returns(MockDbCommand.Object);
+                MockDbCommand.SetupAllProperties();
+                MockDbCommand.Setup(cmd => cmd.ExecuteReader()).Returns(MockDataReader.Object);
+            }
+
+            public Mock<IDbConnection> MockDbConnection { get; private set; }
+            public Mock<IDbCommand> MockDbCommand { get; private set; }
+            public Mock<IDataReader> MockDataReader { get; private set; }
+
+            public IDbConnection CreateConnection()
+            {
+                return MockDbConnection.Object;
+            }
+
+            public virtual DbParameter CreateParameter()
+            {
+                return new Mock<DbParameter>().SetupAllProperties().Object;
+            }
+        }
+    }
+}