fix: Datadog integration not respecting use_custom_source (#6113)

Author: khvn26

api/audit/signals.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Callable, Literal, Protocol, Type
+from typing import TYPE_CHECKING, Any, Callable, Literal, Protocol, Type
 
 from django.conf import settings
 from django.db.models.signals import post_save
@@ -23,6 +23,10 @@
 from webhooks.tasks import call_organisation_webhooks
 from webhooks.webhooks import WebhookEventType
 
+if TYPE_CHECKING:
+    from integrations.datadog.models import DataDogConfiguration
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -129,14 +133,17 @@ def _track_event_async(instance, integration_client):  # type: ignore[no-untyped
 @receiver(post_save, sender=AuditLog)
 @track_only_feature_related_events
 def send_audit_log_event_to_datadog(sender, instance, **kwargs):  # type: ignore[no-untyped-def]
-    data_dog_config = _get_integration_config(instance, "data_dog_config")
+    data_dog_config: "DataDogConfiguration | None" = _get_integration_config(
+        instance, "data_dog_config"
+    )
 
     if not data_dog_config:
         return
 
     data_dog = DataDogWrapper(
         base_url=data_dog_config.base_url,  # type: ignore[arg-type]
         api_key=data_dog_config.api_key,
+        use_custom_source=data_dog_config.use_custom_source,
     )
     _track_event_async(instance, data_dog)  # type: ignore[no-untyped-call]
 

api/tests/unit/audit/test_unit_audit_models.py

@@ -1,3 +1,5 @@
+import json
+
 import pytest
 from pytest_mock import MockerFixture
 
@@ -79,37 +81,81 @@ def test_data_dog_track_event_not_called_on_audit_log_saved_when_wrong(mocker, p
     datadog_mock.track_event_async.assert_not_called()
 
 
-def test_data_dog_track_event_called_on_audit_log_saved_when_correct_type(  # type: ignore[no-untyped-def]
-    project, mocker
-):
+@pytest.mark.parametrize(
+    "use_custom_source, expected_additional_data",
+    [(False, {}), (True, {"source_type_name": "flagsmith"})],
+)
+def test_data_dog_track_event_called_on_audit_log_saved_when_correct_type(
+    project: Project,
+    mocker: MockerFixture,
+    use_custom_source: bool,
+    expected_additional_data: dict[str, str],
+) -> None:
     # Given project configured for Datadog integration
-    datadog_mock = mocker.patch(
-        "integrations.datadog.datadog.DataDogWrapper.track_event_async"
-    )
+    requests_session_mock = mocker.patch(
+        "integrations.datadog.datadog.requests.Session"
+    ).return_value
 
     DataDogConfiguration.objects.create(
-        project=project, base_url='http"//test.com', api_key="123key"
+        project=project,
+        base_url="http://test.com",
+        api_key="123key",
+        use_custom_source=use_custom_source,
     )
 
     # When Audit logs created with correct type
     AuditLog.objects.create(
         project=project,
-        log="Some audit log",
+        log="Some audit log for feature",
         related_object_type=RelatedObjectType.FEATURE.name,
     )
     AuditLog.objects.create(
         project=project,
-        log="Some audit log",
+        log="Some audit log for feature state",
         related_object_type=RelatedObjectType.FEATURE_STATE.name,
     )
     AuditLog.objects.create(
         project=project,
-        log="Some audit log",
+        log="Some audit log for segment",
         related_object_type=RelatedObjectType.SEGMENT.name,
     )
 
     # Then datadog track even triggered for each AuditLog
-    assert 3 == datadog_mock.call_count
+    assert requests_session_mock.post.call_args_list == [
+        mocker.call(
+            "http://test.com/api/v1/events?api_key=123key",
+            data=json.dumps(
+                {
+                    "text": "Some audit log for feature by user system",
+                    "title": "Flagsmith Feature Flag Event",
+                    "tags": ["env:unknown"],
+                    **expected_additional_data,
+                }
+            ),
+        ),
+        mocker.call(
+            "http://test.com/api/v1/events?api_key=123key",
+            data=json.dumps(
+                {
+                    "text": "Some audit log for feature state by user system",
+                    "title": "Flagsmith Feature Flag Event",
+                    "tags": ["env:unknown"],
+                    **expected_additional_data,
+                }
+            ),
+        ),
+        mocker.call(
+            "http://test.com/api/v1/events?api_key=123key",
+            data=json.dumps(
+                {
+                    "text": "Some audit log for segment by user system",
+                    "title": "Flagsmith Feature Flag Event",
+                    "tags": ["env:unknown"],
+                    **expected_additional_data,
+                }
+            ),
+        ),
+    ]
 
 
 def test_audit_log_get_history_record_model_class(mocker):  # type: ignore[no-untyped-def]