greptile comment

Author: parinporecha

examples/celery_integration.py

@@ -21,7 +21,7 @@
 from typing import Any, Optional
 
 from celery import Celery
-from celery.signals import worker_process_init, worker_process_shutdown
+from celery.signals import worker_process_shutdown
 
 import posthog
 from posthog.integrations.celery import PosthogCeleryIntegration

posthog/test/integrations/test_celery_integration.py

@@ -25,21 +25,45 @@ def disconnect(self, handler):
 
 
 class TestPosthogCeleryIntegration(unittest.TestCase):
-    def test_instrument_is_idempotent(self):
+    SIGNAL_NAMES = (
+        "task_prerun",
+        "task_success",
+        "task_failure",
+        "task_retry",
+        "before_task_publish",
+        "after_task_publish",
+        "worker_process_shutdown",
+    )
+
+    def _build_fake_celery(self):
         fake_signals = SimpleNamespace(
-            task_prerun=FakeSignal(),
-            task_success=FakeSignal(),
-            task_failure=FakeSignal(),
-            task_retry=FakeSignal(),
-            before_task_publish=FakeSignal(),
-            after_task_publish=FakeSignal(),
-            worker_process_shutdown=FakeSignal(),
+            **{signal_name: FakeSignal() for signal_name in self.SIGNAL_NAMES}
         )
-
-        integration = PosthogCeleryIntegration()
         fake_celery = ModuleType("celery")
         fake_celery.signals = fake_signals
         fake_celery.__version__ = "5.0.0"
+        return fake_signals, fake_celery
+
+    def _assert_signal_counts(
+        self, fake_signals, expected_connected, expected_disconnected=None
+    ):
+        for signal_name in self.SIGNAL_NAMES:
+            self.assertEqual(
+                len(getattr(fake_signals, signal_name).connected),
+                expected_connected,
+                f"{signal_name} connected count mismatch",
+            )
+            if expected_disconnected is not None:
+                self.assertEqual(
+                    len(getattr(fake_signals, signal_name).disconnected),
+                    expected_disconnected,
+                    f"{signal_name} disconnected count mismatch",
+                )
+
+    def test_instrument_is_idempotent(self):
+        fake_signals, fake_celery = self._build_fake_celery()
+
+        integration = PosthogCeleryIntegration()
 
         with (
             patch.dict("sys.modules", {"celery": fake_celery}),
@@ -48,43 +72,18 @@ def test_instrument_is_idempotent(self):
             integration.instrument()
             integration.instrument()
 
-        for sig in [
-            "task_prerun",
-            "task_success",
-            "task_failure",
-            "task_retry",
-            "before_task_publish",
-            "after_task_publish",
-            "worker_process_shutdown",
-        ]:
-            self.assertEqual(
-                len(getattr(fake_signals, sig).connected),
-                1,
-                f"{sig} connected multiple times",
-            )
+        self._assert_signal_counts(fake_signals, expected_connected=1)
         mock_register.assert_called_once_with(integration.shutdown)
         self.assertEqual(
             fake_signals.worker_process_shutdown.connected[0][0],
             integration._on_worker_process_shutdown,
         )
 
     def test_instrument_and_uninstrument_connect_signals(self):
-        fake_signals = SimpleNamespace(
-            task_prerun=FakeSignal(),
-            task_success=FakeSignal(),
-            task_failure=FakeSignal(),
-            task_retry=FakeSignal(),
-            before_task_publish=FakeSignal(),
-            after_task_publish=FakeSignal(),
-            worker_process_shutdown=FakeSignal(),
-        )
+        fake_signals, fake_celery = self._build_fake_celery()
 
         integration = PosthogCeleryIntegration()
 
-        fake_celery = ModuleType("celery")
-        fake_celery.signals = fake_signals
-        fake_celery.__version__ = "5.0.0"
-
         with (
             patch.dict("sys.modules", {"celery": fake_celery}),
             patch("posthog.integrations.celery.atexit.register"),
@@ -93,23 +92,9 @@ def test_instrument_and_uninstrument_connect_signals(self):
             integration.instrument()
             integration.uninstrument()
 
-        for sig in [
-            "task_prerun",
-            "task_success",
-            "task_failure",
-            "task_retry",
-            "before_task_publish",
-            "after_task_publish",
-            "worker_process_shutdown",
-        ]:
-            self.assertEqual(
-                len(getattr(fake_signals, sig).connected), 1, f"{sig} not connected"
-            )
-            self.assertEqual(
-                len(getattr(fake_signals, sig).disconnected),
-                1,
-                f"{sig} not disconnected",
-            )
+        self._assert_signal_counts(
+            fake_signals, expected_connected=1, expected_disconnected=1
+        )
         mock_unregister.assert_called_once_with(integration.shutdown)
         self.assertEqual(
             fake_signals.worker_process_shutdown.disconnected[0],
@@ -162,22 +147,11 @@ def test_shutdown_is_idempotent(self):
         mock_client.flush.assert_called_once_with()
 
     def test_shutdown_keeps_atexit_registration_when_flush_fails(self):
-        fake_signals = SimpleNamespace(
-            task_prerun=FakeSignal(),
-            task_success=FakeSignal(),
-            task_failure=FakeSignal(),
-            task_retry=FakeSignal(),
-            before_task_publish=FakeSignal(),
-            after_task_publish=FakeSignal(),
-            worker_process_shutdown=FakeSignal(),
-        )
+        fake_signals, fake_celery = self._build_fake_celery()
 
         mock_client = Mock()
         mock_client.flush.side_effect = RuntimeError("flush failed")
         integration = PosthogCeleryIntegration(client=mock_client)
-        fake_celery = ModuleType("celery")
-        fake_celery.signals = fake_signals
-        fake_celery.__version__ = "5.0.0"
 
         with (
             patch.dict("sys.modules", {"celery": fake_celery}),
@@ -194,21 +168,10 @@ def test_shutdown_keeps_atexit_registration_when_flush_fails(self):
         self.assertEqual(len(fake_signals.worker_process_shutdown.disconnected), 1)
 
     def test_reinstrument_after_shutdown_allows_shutdown_again(self):
-        fake_signals = SimpleNamespace(
-            task_prerun=FakeSignal(),
-            task_success=FakeSignal(),
-            task_failure=FakeSignal(),
-            task_retry=FakeSignal(),
-            before_task_publish=FakeSignal(),
-            after_task_publish=FakeSignal(),
-            worker_process_shutdown=FakeSignal(),
-        )
+        fake_signals, fake_celery = self._build_fake_celery()
 
         mock_client = Mock()
         integration = PosthogCeleryIntegration(client=mock_client)
-        fake_celery = ModuleType("celery")
-        fake_celery.signals = fake_signals
-        fake_celery.__version__ = "5.0.0"
 
         with (
             patch.dict("sys.modules", {"celery": fake_celery}),
@@ -223,23 +186,9 @@ def test_reinstrument_after_shutdown_allows_shutdown_again(self):
         self.assertEqual(mock_client.flush.call_count, 2)
         self.assertEqual(mock_register.call_count, 2)
         self.assertEqual(mock_unregister.call_count, 2)
-        for sig in [
-            "task_prerun",
-            "task_success",
-            "task_failure",
-            "task_retry",
-            "before_task_publish",
-            "after_task_publish",
-            "worker_process_shutdown",
-        ]:
-            self.assertEqual(
-                len(getattr(fake_signals, sig).connected), 2, f"{sig} not reconnected"
-            )
-            self.assertEqual(
-                len(getattr(fake_signals, sig).disconnected),
-                2,
-                f"{sig} not disconnected twice",
-            )
+        self._assert_signal_counts(
+            fake_signals, expected_connected=2, expected_disconnected=2
+        )
 
     def test_worker_process_shutdown_hook_calls_shutdown(self):
         integration = PosthogCeleryIntegration(client=Mock())