From 025fba9447769cc71fd8e9f82b5130faa9e2666c Mon Sep 17 00:00:00 2001
From: Omkar Moghe <theomkarmoghe@gmail.com>
Date: Mon, 24 Jun 2024 10:30:40 -0700
Subject: [PATCH] reset context correctly

---
 CHANGELOG.md               | 15 ++++++++-------
 Gemfile.lock               |  2 +-
 lib/lab_coat/experiment.rb |  8 ++++----
 lib/lab_coat/version.rb    |  2 +-
 test/test_experiment.rb    | 21 ++++++++++++++++++++-
 5 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 43f5868..232b932 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,24 +1,25 @@
 ## [0.1.7] - Unreleased
+- Fixes bug where disabled `Experiments` would not reset the runtime context.
 
 ## [0.1.6] - 2024-06-17
-- [#1](https://github.com/omkarmoghe/lab_coat/issues/1)
+- Adds [#1](https://github.com/omkarmoghe/lab_coat/issues/1)
 
 ## [0.1.5] - 2024-05-20
-- Adds `select_observation` to allow users to control which observation value is returned by the experiment. This helps with controlled rollout.
+- Adds `select_observation` to allow users to control which `Observation` value is returned by the `Experiment`. This helps with controlled rollout.
 
 ## [0.1.4] - 2024-04-19
-- Remove the arity check, it's not very intuitive
+- Removes the arity check, it's not very intuitive
 - Adds a `@context` that gets set at runtime and reset after each run. This is a much simpler way for methods to access a shared runtime context that can be set per `run!`.
 
 ## [0.1.3] - 2024-04-17
-- `Experiment` now enforces arity at runtime for the `#enabled?`, `control`, and `candidate` methods.
+- Add an arity check that `Experiment` now enforces at runtime for the `#enabled?`, `control`, and `candidate` methods.
 
 ## [0.1.2] - 2024-04-15
-- use `Benchmark` to capture the duration with more details
-- add `to_h` methods to `Result` and `Observation` for convenience
+- Adds `Benchmark` to capture the duration with more details.
+- Adds `to_h` methods to `Result` and `Observation` for convenience.
 
 ## [0.1.1] - 2024-04-08
-- add `#slug` method to `Observation`
+- Adds `#slug` method to `Observation`.
 
 ## [0.1.0] - 2024-04-08
 - Initial release
diff --git a/Gemfile.lock b/Gemfile.lock
index 58d09ba..46f64ce 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    lab_coat (0.1.6)
+    lab_coat (0.1.7)
 
 GEM
   remote: https://rubygems.org/
diff --git a/lib/lab_coat/experiment.rb b/lib/lab_coat/experiment.rb
index ed77c6f..f88c6d2 100644
--- a/lib/lab_coat/experiment.rb
+++ b/lib/lab_coat/experiment.rb
@@ -105,10 +105,10 @@ def run!(**context) # rubocop:disable Metrics/MethodLength, Metrics/AbcSize
       publish!(result)
 
       # Return the selected observations, control by default.
-      select_observation(result).value.tap do
-        # Reset the context for this run. Done here so that `select_observation` has access to the runtime context.
-        @context = {}
-      end
+      select_observation(result).value
+    ensure
+      # Reset the runtime context before exiting, in all scenarios.
+      @context = {}
     end
   end
 end
diff --git a/lib/lab_coat/version.rb b/lib/lab_coat/version.rb
index e61d97f..790dfbf 100644
--- a/lib/lab_coat/version.rb
+++ b/lib/lab_coat/version.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LabCoat
-  VERSION = "0.1.6"
+  VERSION = "0.1.7"
 end
diff --git a/test/test_experiment.rb b/test/test_experiment.rb
index 7b732b5..a2b7597 100644
--- a/test/test_experiment.rb
+++ b/test/test_experiment.rb
@@ -122,14 +122,33 @@ def test_select_observation # rubocop:disable Metrics/MethodLength
     )
   end
 
-  def test_run!
+  def test_run! # rubocop:disable Metrics/MethodLength
     assert_equal(
       { result: "abc", status: :ok },
       @experiment.run!(num: 1)
     )
+
     assert_equal(
       { result: "abc", status: :ok },
       @experiment.run!(num: 2)
     )
+
+    assert_equal(
+      {},
+      @experiment.context
+    )
+  end
+
+  # Context should be reset if the experiment is disabled and short circuits.
+  def test_short_circuit_run!
+    assert_equal(
+      { result: "abc", status: :ok },
+      @experiment.run!(num: 1)
+    )
+
+    assert_equal(
+      {},
+      @experiment.context
+    )
   end
 end