From 41816695c7aaca83f300dc063f8fa1de29127299 Mon Sep 17 00:00:00 2001
From: phvalguima <pedro.guimaraes@canonical.com>
Date: Wed, 16 Oct 2024 21:19:20 +0200
Subject: [PATCH] [DPE-5677] Fix the append scenarios in replace() (#484)

Update the `replace()` method to (1) be more "greedy" when searching for
matches in the text, including multi-line matches; and (2) fix the write
back to the file.

Currently, it is possible that, if we have a smaller size than the
original file size, we will end up writing:
<new content><left over bytes> -> and this file still has the same size
as it original.

This PR simplifies the logic to decide how to write the changed content.

Closes #483
---
 .../opensearch/v0/helper_conf_setter.py       | 32 +++++----
 tests/unit/lib/test_helper_conf_setter.py     | 65 +++++++++++++++++++
 2 files changed, 80 insertions(+), 17 deletions(-)

diff --git a/lib/charms/opensearch/v0/helper_conf_setter.py b/lib/charms/opensearch/v0/helper_conf_setter.py
index be2f55fc8..12876f8d9 100755
--- a/lib/charms/opensearch/v0/helper_conf_setter.py
+++ b/lib/charms/opensearch/v0/helper_conf_setter.py
@@ -279,23 +279,21 @@ def replace(
         with open(path, "r+") as f:
             data = f.read()
 
-            if regex and old_val and re.compile(old_val).match(data):
-                data = re.sub(r"{}".format(old_val), f"{new_val}", data)
-            elif old_val and old_val in data:
-                data = data.replace(old_val, new_val)
-            elif add_line_if_missing:
-                data += f"{data.rstrip()}\n{new_val}\n"
-
-            if output_type in [OutputType.console, OutputType.all]:
-                logger.info(data)
-
-            if output_type in [OutputType.file, OutputType.all]:
-                if output_file is None or output_file == config_file:
-                    f.seek(0)
-                    f.write(data)
-                else:
-                    with open(output_file, "w") as g:
-                        g.write(data)
+        if regex and old_val and re.compile(old_val, re.MULTILINE).findall(data):
+            data = re.sub(r"{}".format(old_val), f"{new_val}", data)
+        elif old_val and old_val in data:
+            data = data.replace(old_val, new_val)
+        elif add_line_if_missing:
+            data = f"{data.rstrip()}\n{new_val}\n"
+
+        if output_type in [OutputType.console, OutputType.all]:
+            logger.info(data)
+
+        if output_file is None:
+            output_file = config_file
+
+        with open(output_file, "w") as f:
+            f.write(data)
 
     @override
     def append(
diff --git a/tests/unit/lib/test_helper_conf_setter.py b/tests/unit/lib/test_helper_conf_setter.py
index e86c9cf22..0dd17c7e2 100644
--- a/tests/unit/lib/test_helper_conf_setter.py
+++ b/tests/unit/lib/test_helper_conf_setter.py
@@ -4,9 +4,20 @@
 """Unit test for the helper_conf_setter library."""
 import os
 import unittest
+from unittest.mock import mock_open, patch
 
 from charms.opensearch.v0.helper_conf_setter import YamlConfigSetter
 
+REPLACE_TEST_CONTENT = """simple_key: simple_value
+obj:
+  simple_array:
+    - elt1
+    - elt2
+"""
+
+JVM_OPTIONS = """-Xms1g
+-Xmx1g"""
+
 
 class TestHelperConfSetter(unittest.TestCase):
     def setUp(self) -> None:
@@ -82,6 +93,60 @@ def test_delete(self):
         for elt in complex_array:
             self.assertNotEqual(elt["name"], "name1")
 
+    @patch("builtins.open", new_callable=mock_open, read_data=REPLACE_TEST_CONTENT)
+    def test_replace(self, mock_file):
+        """Test replacing values in a file."""
+        input_file = "tests/unit/resources/test_conf.yaml"
+        output_file = "tests/unit/resources/produced.yaml"
+
+        # Test with smaller file size
+        self.conf.replace(input_file, "simple_key", "test", output_file=output_file)
+        mock_file.assert_called_with(output_file, "w")
+        handle = mock_file()
+        handle.write.assert_called_with(
+            "test: simple_value\nobj:\n  simple_array:\n    - elt1\n    - elt2\n"
+        )
+
+        self.conf.replace(input_file, "elt2", "replaced_elt", output_file=output_file)
+        handle.write.assert_called_with(
+            "simple_key: simple_value\nobj:\n  simple_array:\n    - elt1\n    - replaced_elt\n"
+        )
+
+        self.conf.replace(
+            input_file,
+            "non_existing",
+            "new_value",
+            add_line_if_missing=True,
+            output_file=output_file,
+        )
+        handle.write.assert_called_with(
+            "simple_key: simple_value\nobj:\n  simple_array:\n    - elt1\n    - elt2\nnew_value\n"
+        )
+
+    @patch("charms.opensearch.v0.helper_conf_setter.exists")
+    @patch("builtins.open", new_callable=mock_open, read_data=JVM_OPTIONS)
+    def test_multiline_replace(self, mock_file, mock_exists):
+        mock_exists.return_value = True
+
+        self.conf.replace(
+            "jvm.options",
+            "-Xms[0-9]+[kmgKMG]",
+            "-Xms7680k",
+            regex=True,
+        )
+        self.conf.replace(
+            "jvm.options",
+            "-Xmx[0-9]+[kmgKMG]",
+            "-Xmx7680k",
+            regex=True,
+        )
+
+        mock_file.assert_called_with("jvm.options", "w")
+        handle = mock_file()
+
+        handle.write.assert_any_call("-Xms7680k\n-Xmx1g")
+        handle.write.assert_any_call("-Xms1g\n-Xmx7680k")
+
     def tearDown(self) -> None:
         """Cleanup."""
         output = "tests/unit/resources/produced.yaml"