From 9cbc01bcbcd5cbc60d6c801682898c56d669dc91 Mon Sep 17 00:00:00 2001
From: Pavel Krymets <pavel@krymets.com>
Date: Fri, 24 Apr 2020 09:25:47 -0700
Subject: [PATCH] Add guardrails to TestEnvironment (#11570)

---
 .../Azure.Core/tests/RecordedTestBaseTests.cs | 73 +++++++++++++++++++
 .../RecordedTestBase{TEnvironment}.cs         |  4 +-
 .../TestFramework/TestEnvironment.Recorded.cs | 19 +++--
 3 files changed, 90 insertions(+), 6 deletions(-)
 create mode 100644 sdk/core/Azure.Core/tests/RecordedTestBaseTests.cs

diff --git a/sdk/core/Azure.Core/tests/RecordedTestBaseTests.cs b/sdk/core/Azure.Core/tests/RecordedTestBaseTests.cs
new file mode 100644
index 000000000000..af22cbca2046
--- /dev/null
+++ b/sdk/core/Azure.Core/tests/RecordedTestBaseTests.cs
@@ -0,0 +1,73 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using Azure.Core.Testing;
+using NUnit.Framework;
+
+namespace Azure.Core.Tests
+{
+    public class RecordedTestBaseTests
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            Environment.SetEnvironmentVariable("_MOCK_RECORDED", "1");
+            Environment.SetEnvironmentVariable("_MOCK_NOTRECORDED", "2");
+        }
+
+        [Theory]
+        [TestCase(RecordedTestMode.Live)]
+        [TestCase(RecordedTestMode.Playback)]
+        [TestCase(RecordedTestMode.Record)]
+        public void ReadingRecordedValueInCtorThrows(RecordedTestMode mode)
+        {
+            Assert.Throws<InvalidOperationException>(() => new RecordedVariableMisuse(true, mode));
+
+        }
+
+        [Theory]
+        [TestCase(RecordedTestMode.Live)]
+        [TestCase(RecordedTestMode.Playback)]
+        [TestCase(RecordedTestMode.Record)]
+        [TestCase(RecordedTestMode.None)]
+        public void ReadingNonRecordedValueInCtorWorks(RecordedTestMode mode)
+        {
+            var test = new RecordedVariableMisuse(false, mode);
+            Assert.AreEqual("2", test.Value);
+        }
+
+        private class RecordedVariableMisuse : RecordedTestBase<MockTestEnvironment>
+        {
+            // To make NUnit happy
+            public RecordedVariableMisuse(bool isAsync) : base(isAsync)
+            {
+            }
+
+            public RecordedVariableMisuse(bool recorded, RecordedTestMode mode) : base(true, mode)
+            {
+                if (recorded)
+                {
+                    Value = TestEnvironment.RecordedValue;
+                }
+                else
+                {
+                    Value = TestEnvironment.NotRecordedValue;
+                }
+            }
+
+            public string Value { get; }
+        }
+
+        private class MockTestEnvironment: TestEnvironment
+        {
+            public MockTestEnvironment(): base("_mock")
+            {
+            }
+
+            public string RecordedValue => GetRecordedVariable("RECORDED");
+            public string NotRecordedValue => GetVariable("NOTRECORDED");
+
+        }
+    }
+}
\ No newline at end of file
diff --git a/sdk/core/Azure.Core/tests/TestFramework/RecordedTestBase{TEnvironment}.cs b/sdk/core/Azure.Core/tests/TestFramework/RecordedTestBase{TEnvironment}.cs
index 38bb59624d4a..2dc70770773a 100644
--- a/sdk/core/Azure.Core/tests/TestFramework/RecordedTestBase{TEnvironment}.cs
+++ b/sdk/core/Azure.Core/tests/TestFramework/RecordedTestBase{TEnvironment}.cs
@@ -8,17 +8,19 @@ namespace Azure.Core.Testing
         protected RecordedTestBase(bool isAsync) : base(isAsync)
         {
             TestEnvironment = new TEnvironment();
+            TestEnvironment.SetMode(Mode);
         }
 
         protected RecordedTestBase(bool isAsync, RecordedTestMode mode) : base(isAsync, mode)
         {
             TestEnvironment = new TEnvironment();
+            TestEnvironment.SetMode(Mode);
         }
 
         public override void StartTestRecording()
         {
             base.StartTestRecording();
-            TestEnvironment.SetRecording(Recording, Mode == RecordedTestMode.Playback);
+            TestEnvironment.SetRecording(Recording);
         }
 
         public TEnvironment TestEnvironment { get; }
diff --git a/sdk/core/Azure.Core/tests/TestFramework/TestEnvironment.Recorded.cs b/sdk/core/Azure.Core/tests/TestFramework/TestEnvironment.Recorded.cs
index 1f213440c895..caea55db4393 100644
--- a/sdk/core/Azure.Core/tests/TestFramework/TestEnvironment.Recorded.cs
+++ b/sdk/core/Azure.Core/tests/TestFramework/TestEnvironment.Recorded.cs
@@ -11,20 +11,24 @@ namespace Azure.Core.Testing
     public partial class TestEnvironment
     {
         private TestRecording _recording;
-        private bool _playback;
+        private RecordedTestMode _mode;
 
-        public void SetRecording(TestRecording recording, bool playback)
+        public void SetMode(RecordedTestMode mode)
+        {
+            _mode = mode;
+        }
+
+        public void SetRecording(TestRecording recording)
         {
             _credential = null;
             _recording = recording;
-            _playback = playback;
         }
 
         partial void GetRecordedValue(string name, ref string value)
         {
             if (_recording == null)
             {
-                return;
+                throw new InvalidOperationException("Recorded value should not be retrieved outside the test method invocation");
             }
 
             value =  _recording.GetVariable(name, null);
@@ -32,11 +36,16 @@ partial void GetRecordedValue(string name, ref string value)
 
         partial void GetIsPlayback(ref bool playback)
         {
-            playback = _playback;
+            playback = _mode == RecordedTestMode.Playback;
         }
 
         partial void SetRecordedValue(string name, string value)
         {
+            if (_recording == null)
+            {
+                throw new InvalidOperationException("Recorded value should not be retrieved outside the test method invocation");
+            }
+
             _recording?.SetVariable(name, value);
         }
     }