From 2887886527edf0eb69ffcbf7a43002dad3b9730c Mon Sep 17 00:00:00 2001
From: Sherzod Karimov <sherzod.karimov@datadoghq.com>
Date: Thu, 2 May 2024 14:09:56 -0400
Subject: [PATCH] refactor and add tests

---
 datadog_sync/model/logs_indexes_order.py      | 23 +++++++-----
 .../resources/test_logs_indexes_order.py      | 36 +++++++++++++++++++
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/datadog_sync/model/logs_indexes_order.py b/datadog_sync/model/logs_indexes_order.py
index 7824d477..177f571b 100644
--- a/datadog_sync/model/logs_indexes_order.py
+++ b/datadog_sync/model/logs_indexes_order.py
@@ -57,14 +57,7 @@ async def create_resource(self, _id: str, resource: Dict) -> Tuple[str, Dict]:
 
     async def update_resource(self, _id: str, resource: Dict) -> Tuple[str, Dict]:
         destination_resources = self.destination_indexes_order or self.resource_config.destination_resources[_id]
-        ids_to_omit = set(resource["index_names"]) - set(destination_resources["index_names"])
-
-        extra_ids_to_include = [
-            _id for _id in destination_resources["index_names"] if _id not in resource["index_names"]
-        ]
-
-        resource["index_names"] = [_id for _id in resource["index_names"] if _id not in ids_to_omit]
-        resource["index_names"] = resource["index_names"] + extra_ids_to_include
+        self.handle_additional_indexes(resource, destination_resources)
 
         destination_client = self.config.destination_client
         resp = await destination_client.put(self.resource_config.base_path, resource)
@@ -91,3 +84,17 @@ async def get_destination_indexes_order(self):
         resp = await self.get_resources(destination_client)
 
         return resp[0]
+
+    @staticmethod
+    def handle_additional_indexes(resource, destination_resource) -> None:
+        # Logs index order requires all logs indexes in the destination org to be included in the payload
+        # Additional indexes in the source org which need to be removed from the payload
+        ids_to_omit = set(resource["index_names"]) - set(destination_resource["index_names"])
+        resource["index_names"] = [_id for _id in resource["index_names"] if _id not in ids_to_omit]
+
+        # Add back additional indexes present in the destination org while retaining the relative ordering
+        # of the additional indexes
+        extra_ids_to_include = [
+            _id for _id in destination_resource["index_names"] if _id not in resource["index_names"]
+        ]
+        resource["index_names"] = resource["index_names"] + extra_ids_to_include
diff --git a/tests/integration/resources/test_logs_indexes_order.py b/tests/integration/resources/test_logs_indexes_order.py
index 05326aaf..e30ea27a 100644
--- a/tests/integration/resources/test_logs_indexes_order.py
+++ b/tests/integration/resources/test_logs_indexes_order.py
@@ -16,3 +16,39 @@ class TestLogsIndexesOrder(BaseResourcesTestClass):
     @pytest.mark.skip(reason="resource is only updated by default")
     def test_resource_update_sync(self):
         pass
+
+
+@pytest.mark.parametrize(
+    "resource, destination_resource, expected",
+    [
+        (
+            {"index_names": ["index1", "index2", "index3"]},
+            {"index_names": ["index3", "index2", "index4"]},
+            {"index_names": ["index2", "index3", "index4"]},
+        ),
+        (
+            {"index_names": ["index1"]},
+            {"index_names": ["index3", "index1", "index4"]},
+            {"index_names": ["index1", "index3", "index4"]},
+        ),
+        (
+            {"index_names": ["index1", "index2", "index3"]},
+            {"index_names": ["index3", "index1"]},
+            {"index_names": ["index1", "index3"]},
+        ),
+        (
+            {"index_names": ["index1", "index2", "index3"]},
+            {"index_names": ["index1"]},
+            {"index_names": ["index1"]},
+        ),
+        (
+            {"index_names": ["index1"]},
+            {"index_names": ["index5", "index1", "index3", "index4"]},
+            {"index_names": ["index1", "index5", "index3", "index4"]},
+        ),
+    ],
+)
+def test_handle_index_diff(resource, destination_resource, expected):
+    LogsIndexesOrder.handle_additional_indexes(resource, destination_resource)
+
+    assert resource == expected