Fix flaky SetupSequence tests (#7051)

Author: amadeuszl

test/Libraries/Microsoft.Extensions.Diagnostics.ResourceMonitoring.Tests/Windows/WindowsContainerSnapshotProviderTests.cs

@@ -198,12 +198,13 @@ public void SnapshotProvider_EmitsCpuTimeMetric()
         updatedAccountingInfo.TotalKernelTime = 2500;
         updatedAccountingInfo.TotalUserTime = 1500;
 
-        _jobHandleMock.SetupSequence(j => j.GetBasicAccountingInfo())
-            .Returns(_accountingInfo)
-            .Returns(_accountingInfo)
-            .Returns(updatedAccountingInfo)
-            .Returns(updatedAccountingInfo)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
+        var callCount = 0;
+        _jobHandleMock.Setup(j => j.GetBasicAccountingInfo())
+            .Returns(() =>
+            {
+                callCount++;
+                return callCount <= 2 ? _accountingInfo : updatedAccountingInfo;
+            });
 
         _sysInfo.NumberOfProcessors = 2;
 
@@ -231,14 +232,14 @@ public void SnapshotProvider_EmitsCpuTimeMetric()
         var snapshot = metricCollector.GetMeasurementSnapshot();
         Assert.Equal(2, snapshot.Count);
         Assert.Contains(_accountingInfo.TotalKernelTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
-        Assert.Contains(_accountingInfo.TotalKernelTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
+        Assert.Contains(_accountingInfo.TotalUserTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
 
         // Step #1 - simulate 1 millisecond passing and collect metrics again:
         fakeClock.Advance(TimeSpan.FromMilliseconds(1));
         metricCollector.RecordObservableInstruments();
         snapshot = metricCollector.GetMeasurementSnapshot();
         Assert.Contains(updatedAccountingInfo.TotalKernelTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
-        Assert.Contains(updatedAccountingInfo.TotalKernelTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
+        Assert.Contains(updatedAccountingInfo.TotalUserTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
 
         // Step #2 - simulate 1 millisecond passing and collect metrics again:
         fakeClock.Advance(TimeSpan.FromMilliseconds(1));
@@ -247,7 +248,7 @@ public void SnapshotProvider_EmitsCpuTimeMetric()
 
         // CPU time should be the same as before, as we're not simulating any CPU usage:
         Assert.Contains(updatedAccountingInfo.TotalKernelTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
-        Assert.Contains(updatedAccountingInfo.TotalKernelTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
+        Assert.Contains(updatedAccountingInfo.TotalUserTime / (double)TimeSpan.TicksPerSecond, snapshot.Select(m => m.Value));
     }
 
     [Theory]
@@ -262,12 +263,13 @@ public void SnapshotProvider_EmitsCpuMetrics(string instrumentName, bool useZero
         updatedAccountingInfo.TotalKernelTime = 2500;
         updatedAccountingInfo.TotalUserTime = 1500;
 
-        _jobHandleMock.SetupSequence(j => j.GetBasicAccountingInfo())
-            .Returns(() => _accountingInfo)
-            .Returns(updatedAccountingInfo)
-            .Returns(updatedAccountingInfo)
-            .Returns(updatedAccountingInfo)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
+        var callCount = 0;
+        _jobHandleMock.Setup(j => j.GetBasicAccountingInfo())
+            .Returns(() =>
+            {
+                callCount++;
+                return callCount == 1 ? _accountingInfo : updatedAccountingInfo;
+            });
 
         _sysInfo.NumberOfProcessors = 2;
 
@@ -330,15 +332,21 @@ public void SnapshotProvider_EmitsMemoryMetrics(string instrumentName, bool useZ
         _appMemoryUsage = 200UL;
         ulong updatedAppMemoryUsage = 600UL;
 
-        _processInfoMock.SetupSequence(p => p.GetCurrentProcessMemoryUsage())
-            .Returns(() => _appMemoryUsage)
-            .Returns(updatedAppMemoryUsage)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
-
-        _processInfoMock.SetupSequence(p => p.GetMemoryUsage())
-            .Returns(() => _appMemoryUsage)
-            .Returns(updatedAppMemoryUsage)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
+        var processMemoryCallCount = 0;
+        _processInfoMock.Setup(p => p.GetCurrentProcessMemoryUsage())
+            .Returns(() =>
+            {
+                processMemoryCallCount++;
+                return processMemoryCallCount <= 1 ? _appMemoryUsage : updatedAppMemoryUsage;
+            });
+
+        var containerMemoryCallCount = 0;
+        _processInfoMock.Setup(p => p.GetMemoryUsage())
+            .Returns(() =>
+            {
+                containerMemoryCallCount++;
+                return containerMemoryCallCount <= 1 ? _appMemoryUsage : updatedAppMemoryUsage;
+            });
 
         var fakeClock = new FakeTimeProvider();
         using var meter = new Meter(nameof(SnapshotProvider_EmitsMemoryMetrics));
@@ -387,15 +395,21 @@ public void SnapshotProvider_TestMemoryMetricsTogether()
         ulong updatedAppMemoryUsage = 600UL;
         ulong updatedContainerMemoryUsage = 1200UL;
 
-        _processInfoMock.SetupSequence(p => p.GetCurrentProcessMemoryUsage())
-            .Returns(() => _appMemoryUsage)
-            .Returns(updatedAppMemoryUsage)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
-
-        _processInfoMock.SetupSequence(p => p.GetMemoryUsage())
-            .Returns(() => containerMemoryUsage)
-            .Returns(updatedContainerMemoryUsage)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
+        var processMemoryCallCount = 0;
+        _processInfoMock.Setup(p => p.GetCurrentProcessMemoryUsage())
+            .Returns(() =>
+            {
+                processMemoryCallCount++;
+                return processMemoryCallCount <= 1 ? _appMemoryUsage : updatedAppMemoryUsage;
+            });
+
+        var containerMemoryCallCount = 0;
+        _processInfoMock.Setup(p => p.GetMemoryUsage())
+            .Returns(() =>
+            {
+                containerMemoryCallCount++;
+                return containerMemoryCallCount <= 1 ? containerMemoryUsage : updatedContainerMemoryUsage;
+            });
 
         var fakeClock = new FakeTimeProvider();
         using var meter = new Meter(nameof(SnapshotProvider_TestMemoryMetricsTogether));
@@ -465,17 +479,31 @@ public void SnapshotProvider_EmitsMemoryUsageMetric()
         const ulong UpdatedAppMemoryUsage = 600UL;
         const ulong UpdatedAppMemoryUsage2 = 300UL;
 
-        _processInfoMock.SetupSequence(p => p.GetCurrentProcessMemoryUsage())
-            .Returns(() => _appMemoryUsage)
-            .Returns(UpdatedAppMemoryUsage)
-            .Returns(UpdatedAppMemoryUsage2)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
-
-        _processInfoMock.SetupSequence(p => p.GetMemoryUsage())
-            .Returns(() => _appMemoryUsage)
-            .Returns(UpdatedAppMemoryUsage)
-            .Returns(UpdatedAppMemoryUsage2)
-            .Throws(new InvalidOperationException("We shouldn't hit here..."));
+        var processMemoryCallCount = 0;
+        _processInfoMock.Setup(p => p.GetCurrentProcessMemoryUsage())
+            .Returns(() =>
+            {
+                processMemoryCallCount++;
+                return processMemoryCallCount switch
+                {
+                    1 => _appMemoryUsage,
+                    2 => UpdatedAppMemoryUsage,
+                    _ => UpdatedAppMemoryUsage2
+                };
+            });
+
+        var containerMemoryCallCount = 0;
+        _processInfoMock.Setup(p => p.GetMemoryUsage())
+            .Returns(() =>
+            {
+                containerMemoryCallCount++;
+                return containerMemoryCallCount switch
+                {
+                    1 => _appMemoryUsage,
+                    2 => UpdatedAppMemoryUsage,
+                    _ => UpdatedAppMemoryUsage2
+                };
+            });
 
         var fakeClock = new FakeTimeProvider();
         using var meter = new Meter(nameof(SnapshotProvider_EmitsMemoryMetrics));