From 8ab00028cc4f7e74cb539a8246e0d7dcfd2ec565 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Srdjan=20Jov=C4=8Di=C4=87?= <srdjanj@microsoft.com>
Date: Tue, 19 May 2020 16:25:09 -0700
Subject: [PATCH] Add activity unit tests for TelemetryActivity.

---
 .../TelemetryActivityTests.cs                 | 33 ++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/test/NuGet.Core.Tests/NuGet.Common.Test/TelemetryActivityTests.cs b/test/NuGet.Core.Tests/NuGet.Common.Test/TelemetryActivityTests.cs
index b53ca416e65..c5a871000cb 100644
--- a/test/NuGet.Core.Tests/NuGet.Common.Test/TelemetryActivityTests.cs
+++ b/test/NuGet.Core.Tests/NuGet.Common.Test/TelemetryActivityTests.cs
@@ -13,14 +13,25 @@ namespace NuGet.Common.Test
 {
     public class TelemetryActivityTests
     {
-        private readonly Mock<INuGetTelemetryService> _telemetryService = new Mock<INuGetTelemetryService>();
+        private readonly Mock<INuGetTelemetryService> _telemetryService = new Mock<INuGetTelemetryService>(MockBehavior.Strict);
         private TelemetryEvent _telemetryEvent;
 
+        private readonly Mock<IDisposable> _activity = new Mock<IDisposable>(MockBehavior.Strict);
+        private bool _activityDisposed;
+        private string _activityName;
+
         public TelemetryActivityTests()
         {
             _telemetryService.Setup(x => x.EmitTelemetryEvent(It.IsAny<TelemetryEvent>()))
                 .Callback<TelemetryEvent>(x => _telemetryEvent = x);
 
+            _telemetryService.Setup(x => x.StartActivity(It.IsAny<string>()))
+                .Callback<string>(x => _activityName = x)
+                .Returns(_activity.Object);
+
+            _activity.Setup(x => x.Dispose())
+                     .Callback(() => _activityDisposed = true);
+
             TelemetryActivity.NuGetTelemetryService = _telemetryService.Object;
         }
 
@@ -125,6 +136,26 @@ public void Dispose_Always_EmitsDuration()
             Assert.InRange(duration, 0d, 10d);
         }
 
+        [Fact]
+        public void Create_will_start_activity()
+        {
+            using (var telemetry = TelemetryActivity.Create(CreateNewTelemetryEvent()))
+            {
+            }
+
+            Assert.Equal("testEvent", _activityName);
+        }
+
+        [Fact]
+        public void Dispose_will_dispose_activity()
+        {
+            using (var telemetry = TelemetryActivity.Create(CreateNewTelemetryEvent()))
+            {
+            }
+
+            Assert.True(_activityDisposed);
+        }
+
         private static TelemetryEvent CreateNewTelemetryEvent()
         {
             return new TelemetryEvent("testEvent", new Dictionary<string, object>());