diff --git a/src/LaunchDarkly/Integrations/TestData.php b/src/LaunchDarkly/Integrations/TestData.php
index 718b2c9c..239e4ef5 100644
--- a/src/LaunchDarkly/Integrations/TestData.php
+++ b/src/LaunchDarkly/Integrations/TestData.php
@@ -110,7 +110,7 @@ public function update(FlagBuilder $flagBuilder): TestData
 
         $oldFlag = $this->_currentFlags[$key] ?? null;
         if ($oldFlag) {
-            $oldVersion = $oldFlag['version'];
+            $oldVersion = $oldFlag->getVersion();
         }
 
         $newFlag = $flagBuilder->build($oldVersion + 1);
diff --git a/tests/Integrations/TestDataTest.php b/tests/Integrations/TestDataTest.php
index 1e120a92..066eeb7c 100644
--- a/tests/Integrations/TestDataTest.php
+++ b/tests/Integrations/TestDataTest.php
@@ -430,6 +430,37 @@ public function testCanSetAndGetFeatureFlag()
         $this->assertEquals($expectedFeatureFlag, $featureFlag);
     }
 
+    public function testCanSetAndResetFeatureFlag()
+    {
+        $key = 'test-flag';
+        $expectedUpdatedFlagJson = [
+            'key' => $key,
+            'version' => 2,
+            'deleted' => false,
+            'on' => true,
+            'targets' => [],
+            'rules' => [],
+            'offVariation' => 1,
+            'fallthrough' => ['variation' => 2],
+            'variations' => ['red', 'amber', 'green'],
+
+            /* Required FeatureFlag fields */
+            'salt' => null,
+            'prerequisites' => [],
+        ];
+        $expectedUpdatedFeatureFlag = FeatureFlag::decode($expectedUpdatedFlagJson);
+
+        $td = new TestData();
+        $flag = $td->flag($key);
+        $td->update($flag);
+        
+        $updatedFlag = $flag->variations('red', 'amber', 'green')->fallthroughVariation(2);
+        $td->update($updatedFlag);
+
+        $featureFlag = $td->getFeature($key);
+        $this->assertEquals($expectedUpdatedFeatureFlag, $featureFlag);
+    }
+
     public function testFlagBuilderCanAddAndBuildRules()
     {
         $td = new TestData();