wip

grafana · Feb 14, 2025 · 45381f8 · 45381f8
1 parent a6e4cac
commit 45381f8
Show file tree

Hide file tree

Showing 4 changed files with 184 additions and 200 deletions.
diff --git a/tools/migrators/lib/tests/pagerduty/test_matching.py b/tools/migrators/lib/tests/pagerduty/test_matching.py
@@ -1330,7 +1330,7 @@
             "scheduled_actions": [],
         },
         "oncall_integration": None,
-        "oncall_type": None,
+        "oncall_type": "amazon_sns",
         "is_escalation_policy_flawed": False,
     },
     {
@@ -1420,7 +1420,7 @@
             "scheduled_actions": [],
         },
         "oncall_integration": None,
-        "oncall_type": None,
+        "oncall_type": "amazon_sns",
         "is_escalation_policy_flawed": True,
     },
     {
@@ -1510,7 +1510,7 @@
             "scheduled_actions": [],
         },
         "oncall_integration": None,
-        "oncall_type": None,
+        "oncall_type": "amazon_sns",
         "is_escalation_policy_flawed": True,
     },
     {

diff --git a/tools/migrators/lib/tests/pagerduty/test_migrate.py b/tools/migrators/lib/tests/pagerduty/test_migrate.py
@@ -1,14 +1,4 @@
-import os
 from unittest.mock import call, patch
-
-
-# Mock environment variables before importing any modules that use them
-os.environ["MIGRATING_FROM"] = "pagerduty"
-os.environ["ONCALL_API_TOKEN"] = "test-token"
-os.environ["ONCALL_API_URL"] = "http://test"
-os.environ["PAGERDUTY_API_TOKEN"] = "test-pd-token"
-os.environ["MODE"] = "plan"
-
 from lib.pagerduty.migrate import migrate, filter_schedules, filter_escalation_policies, filter_integrations
 
 

diff --git a/tools/migrators/lib/tests/pagerduty/test_migrate_notification_rules.py b/tools/migrators/lib/tests/pagerduty/test_migrate_notification_rules.py
@@ -1,14 +1,129 @@
 from unittest.mock import call, patch
 
-from lib.oncall.api_client import OnCallAPIClient
 from lib.pagerduty.resources.notification_rules import migrate_notification_rules
 
 
-@patch.object(OnCallAPIClient, "delete")
-@patch.object(OnCallAPIClient, "create")
-def test_migrate_notification_rules(api_client_create_mock, api_client_delete_mock):
-    migrate_notification_rules(
-        {
+class TestNotificationRulesPreservation:
+    def setup_method(self):
+        self.pd_user = {
+            "id": "U1",
+            "name": "Test User",
+            "email": "test@example.com",
+            "notification_rules": [{
+                "id": "PD1",
+                "urgency": "high",
+                "start_delay_in_minutes": 0,
+                "contact_method": {
+                    "type": "email_contact_method"
+                }
+            }]
+        }
+        self.oncall_user = {
+            "id": "OC1",
+            "email": "test@example.com",
+            "notification_rules": []
+        }
+        self.pd_user["oncall_user"] = self.oncall_user
+
+    @patch("lib.pagerduty.resources.notification_rules.PRESERVE_EXISTING_USER_NOTIFICATION_RULES", True)
+    @patch("lib.pagerduty.resources.notification_rules.OnCallAPIClient")
+    def test_existing_notification_rules_are_preserved(self, MockOnCallAPIClient):
+        # Setup user with existing notification rules
+        self.oncall_user["notification_rules"] = [{"id": "NR1"}]
+
+        # Run migration
+        migrate_notification_rules(self.pd_user)
+
+        # Verify no notification rules were migrated
+        MockOnCallAPIClient.create.assert_not_called()
+        MockOnCallAPIClient.delete.assert_not_called()
+
+    @patch("lib.pagerduty.resources.notification_rules.PRESERVE_EXISTING_USER_NOTIFICATION_RULES", True)
+    @patch("lib.pagerduty.resources.notification_rules.OnCallAPIClient")
+    def test_notification_rules_migrated_when_none_exist(self, MockOnCallAPIClient):
+        # Run migration
+        migrate_notification_rules(self.pd_user)
+
+        # Verify notification rules were migrated for both important and non-important cases
+        expected_calls = [
+            call("personal_notification_rules", {
+                "user_id": "OC1",
+                "type": "notify_by_email",
+                "important": False
+            }),
+            call("personal_notification_rules", {
+                "user_id": "OC1",
+                "type": "notify_by_email",
+                "important": True
+            })
+        ]
+        MockOnCallAPIClient.create.assert_has_calls(expected_calls)
+        MockOnCallAPIClient.delete.assert_not_called()
+
+    @patch("lib.pagerduty.resources.notification_rules.PRESERVE_EXISTING_USER_NOTIFICATION_RULES", False)
+    @patch("lib.pagerduty.resources.notification_rules.OnCallAPIClient")
+    def test_existing_notification_rules_are_replaced_when_preserve_is_false(self, MockOnCallAPIClient):
+        # Setup user with existing notification rules
+        self.oncall_user["notification_rules"] = [
+            {"id": "NR1", "important": False},
+            {"id": "NR2", "important": True}
+        ]
+
+        # Run migration
+        migrate_notification_rules(self.pd_user)
+
+        # Verify old rules were deleted
+        expected_delete_calls = [
+            call("personal_notification_rules/NR1"),
+            call("personal_notification_rules/NR2")
+        ]
+        MockOnCallAPIClient.delete.assert_has_calls(expected_delete_calls, any_order=True)
+
+        # Verify new rules were created
+        expected_create_calls = [
+            call("personal_notification_rules", {
+                "user_id": "OC1",
+                "type": "notify_by_email",
+                "important": False
+            }),
+            call("personal_notification_rules", {
+                "user_id": "OC1",
+                "type": "notify_by_email",
+                "important": True
+            })
+        ]
+        MockOnCallAPIClient.create.assert_has_calls(expected_create_calls)
+
+    @patch("lib.pagerduty.resources.notification_rules.PRESERVE_EXISTING_USER_NOTIFICATION_RULES", False)
+    @patch("lib.pagerduty.resources.notification_rules.OnCallAPIClient")
+    def test_notification_rules_migrated_when_none_exist_and_preserve_is_false(self, MockOnCallAPIClient):
+        # Run migration
+        migrate_notification_rules(self.pd_user)
+
+        # Verify no rules were deleted (since none existed)
+        MockOnCallAPIClient.delete.assert_not_called()
+
+        # Verify new rules were created
+        expected_create_calls = [
+            call("personal_notification_rules", {
+                "user_id": "OC1",
+                "type": "notify_by_email",
+                "important": False
+            }),
+            call("personal_notification_rules", {
+                "user_id": "OC1",
+                "type": "notify_by_email",
+                "important": True
+            })
+        ]
+        MockOnCallAPIClient.create.assert_has_calls(expected_create_calls)
+
+    @patch("lib.pagerduty.resources.notification_rules.PRESERVE_EXISTING_USER_NOTIFICATION_RULES", False)
+    @patch("lib.pagerduty.resources.notification_rules.OnCallAPIClient")
+    def test_complex_notification_rules_migration(self, MockOnCallAPIClient):
+        # Test a more complex case with multiple notification methods and delays
+        user = {
+            "email": "test@example.com",
             "notification_rules": [
                 {
                     "contact_method": {"type": "sms_contact_method"},
@@ -29,57 +144,63 @@ def test_migrate_notification_rules(api_client_create_mock, api_client_delete_mo
                 ],
             },
         }
-    )
-
-    assert api_client_create_mock.call_args_list == [
-        call(
-            "personal_notification_rules",
-            {
-                "user_id": "EXISTING_USER_ID",
-                "type": "notify_by_sms",
-                "important": False,
-            },
-        ),
-        call(
-            "personal_notification_rules",
-            {
-                "user_id": "EXISTING_USER_ID",
-                "type": "wait",
-                "duration": 300,
-                "important": False,
-            },
-        ),
-        call(
-            "personal_notification_rules",
-            {
-                "user_id": "EXISTING_USER_ID",
-                "type": "notify_by_mobile_app",
-                "important": False,
-            },
-        ),
-        call(
-            "personal_notification_rules",
-            {"user_id": "EXISTING_USER_ID", "type": "notify_by_sms", "important": True},
-        ),
-        call(
-            "personal_notification_rules",
-            {
-                "user_id": "EXISTING_USER_ID",
-                "type": "wait",
-                "duration": 300,
-                "important": True,
-            },
-        ),
-        call(
-            "personal_notification_rules",
-            {
-                "user_id": "EXISTING_USER_ID",
-                "type": "notify_by_mobile_app",
-                "important": True,
-            },
-        ),
-    ]
-    assert api_client_delete_mock.call_args_list == [
-        call("personal_notification_rules/EXISTING_RULE_ID_1"),
-        call("personal_notification_rules/EXISTING_RULE_ID_2"),
-    ]
+
+        migrate_notification_rules(user)
+
+        # Verify old rules were deleted
+        expected_delete_calls = [
+            call("personal_notification_rules/EXISTING_RULE_ID_1"),
+            call("personal_notification_rules/EXISTING_RULE_ID_2")
+        ]
+        MockOnCallAPIClient.delete.assert_has_calls(expected_delete_calls, any_order=True)
+
+        # Verify new rules were created in correct order with correct delays
+        expected_create_calls = [
+            call(
+                "personal_notification_rules",
+                {
+                    "user_id": "EXISTING_USER_ID",
+                    "type": "notify_by_sms",
+                    "important": False,
+                },
+            ),
+            call(
+                "personal_notification_rules",
+                {
+                    "user_id": "EXISTING_USER_ID",
+                    "type": "wait",
+                    "duration": 300,
+                    "important": False,
+                },
+            ),
+            call(
+                "personal_notification_rules",
+                {
+                    "user_id": "EXISTING_USER_ID",
+                    "type": "notify_by_mobile_app",
+                    "important": False,
+                },
+            ),
+            call(
+                "personal_notification_rules",
+                {"user_id": "EXISTING_USER_ID", "type": "notify_by_sms", "important": True},
+            ),
+            call(
+                "personal_notification_rules",
+                {
+                    "user_id": "EXISTING_USER_ID",
+                    "type": "wait",
+                    "duration": 300,
+                    "important": True,
+                },
+            ),
+            call(
+                "personal_notification_rules",
+                {
+                    "user_id": "EXISTING_USER_ID",
+                    "type": "notify_by_mobile_app",
+                    "important": True,
+                },
+            ),
+        ]
+        MockOnCallAPIClient.create.assert_has_calls(expected_create_calls)