Fix issue #4906 by removing version stages from old versions (#4907)

moto/secretsmanager/models.py

@@ -111,6 +111,12 @@ def reset_default_version(self, secret_version, version_id):
         self.versions[version_id] = secret_version
         self.default_version_id = version_id
 
+    def remove_version_stages_from_old_versions(self, version_stages):
+        for version_stage in version_stages:
+            for old_version in self.versions.values():
+                if version_stage in old_version["version_stages"]:
+                    old_version["version_stages"].remove(version_stage)
+
     def delete(self, deleted_date):
         self.deleted_date = deleted_date
 
@@ -377,6 +383,7 @@ def _add_secret(
             if "AWSCURRENT" in version_stages:
                 secret.reset_default_version(secret_version, version_id)
             else:
+                secret.remove_version_stages_from_old_versions(version_stages)
                 secret.versions[version_id] = secret_version
         else:
             secret = FakeSecret(

tests/test_secretsmanager/test_secretsmanager.py

@@ -1071,6 +1071,54 @@ def test_after_put_secret_value_version_stages_can_get_current():
     assert get_dict["VersionStages"] == ["AWSCURRENT"]
 
 
+@mock_secretsmanager
+def test_after_put_secret_value_version_stages_can_get_current_with_custom_version_stage():
+    conn = boto3.client("secretsmanager", region_name="us-west-2")
+
+    # Creation.
+    first_version_id = "eb41453f-25bb-4025-b7f4-850cfca0ce71"
+    first_secret_string = "first_secret_string"
+    conn.create_secret(
+        Name=DEFAULT_SECRET_NAME,
+        SecretString=first_secret_string,
+        ClientRequestToken=first_version_id,
+    )
+
+    # Use PutSecretValue to push a new version with new version stages.
+    second_version_id = "eb41453f-25bb-4025-b7f4-850cfca0ce72"
+    conn.put_secret_value(
+        SecretId=DEFAULT_SECRET_NAME,
+        SecretString="second_secret_string",
+        VersionStages=["SAMPLESTAGE1", "SAMPLESTAGE0"],
+        ClientRequestToken=second_version_id,
+    )
+    # Create a third version with one of the old stages
+    third_version_id = "eb41453f-25bb-4025-b7f4-850cfca0ce73"
+    third_secret_string = "third_secret_string"
+    conn.put_secret_value(
+        SecretId=DEFAULT_SECRET_NAME,
+        SecretString=third_secret_string,
+        VersionStages=["SAMPLESTAGE1"],
+        ClientRequestToken=third_version_id,
+    )
+
+    # Get current with the stage label of the third version.
+    get_dict = conn.get_secret_value(
+        SecretId=DEFAULT_SECRET_NAME, VersionStage="SAMPLESTAGE1"
+    )
+    versions = conn.list_secret_version_ids(SecretId=DEFAULT_SECRET_NAME)["Versions"]
+    versions_by_key = {version["VersionId"]: version for version in versions}
+    # Check if indeed the third version is returned
+    assert get_dict
+    assert get_dict["VersionId"] == third_version_id
+    assert get_dict["SecretString"] == third_secret_string
+    assert get_dict["VersionStages"] == ["SAMPLESTAGE1"]
+    # Check if all the versions have the proper labels
+    assert versions_by_key[first_version_id]["VersionStages"] == ["AWSCURRENT"]
+    assert versions_by_key[second_version_id]["VersionStages"] == ["SAMPLESTAGE0"]
+    assert versions_by_key[third_version_id]["VersionStages"] == ["SAMPLESTAGE1"]
+
+
 @mock_secretsmanager
 def test_after_put_secret_value_version_stages_pending_can_get_current():
     conn = boto3.client("secretsmanager", region_name="us-west-2")