feat: Set base url for segment (#4684)

api/conftest.py

@@ -18,6 +18,7 @@
 from rest_framework.authtoken.models import Token
 from rest_framework.test import APIClient
 from task_processor.task_run_method import TaskRunMethod
+from urllib3 import HTTPResponse
 from urllib3.connectionpool import HTTPConnectionPool
 from xdist import get_xdist_worker_id
 
@@ -163,7 +164,7 @@ def urlopen_mock(
         url: str,
         *args,
         **kwargs,
-    ) -> HTTPConnectionPool.ResponseCls:
+    ) -> HTTPResponse:
         if self.host in allowed_hosts:
             return original_urlopen(self, method, url, *args, **kwargs)
 

api/integrations/segment/constants.py

@@ -0,0 +1,2 @@
+DEFAULT_BASE_URL = "api.segment.io/"
+DUBLIN_BASE_URL = "events.eu1.segmentapis.com/"

api/integrations/segment/migrations/0005_set_base_url_to_default.py

@@ -0,0 +1,23 @@
+from django.apps.registry import Apps
+from django.db import migrations
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+
+from integrations.segment import constants
+
+
+def set_base_url(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
+    SegmentConfiguration = apps.get_model("segment", "SegmentConfiguration")
+    SegmentConfiguration.objects.filter(base_url__isnull=True).update(
+        base_url=constants.DEFAULT_BASE_URL,
+    )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("segment", "0004_segmentconfiguration_deleted_at"),
+    ]
+
+    operations = [
+        migrations.RunPython(set_base_url, reverse_code=migrations.RunPython.noop),
+    ]

api/integrations/segment/segment.py

@@ -14,12 +14,12 @@
 
 class SegmentWrapper(AbstractBaseIdentityIntegrationWrapper):
     def __init__(self, config: SegmentConfiguration):
-        api_key = config.api_key
-        self.analytics = SegmentClient(write_key=api_key, sync_mode=True)
+        self.analytics = SegmentClient(
+            write_key=config.api_key, sync_mode=True, host=config.base_url
+        )
 
     def _identify_user(self, data: dict) -> None:
         self.analytics.identify(**data)
-        logger.debug("Sent event to Segment.")
 
     def generate_user_data(
         self,

api/integrations/segment/serializers.py

@@ -1,10 +1,23 @@
+from rest_framework import serializers
+
 from integrations.common.serializers import (
     BaseEnvironmentIntegrationModelSerializer,
 )
+from integrations.segment import constants
 from integrations.segment.models import SegmentConfiguration
 
 
 class SegmentConfigurationSerializer(BaseEnvironmentIntegrationModelSerializer):
+
+    base_url = serializers.ChoiceField(
+        choices=[
+            constants.DEFAULT_BASE_URL,
+            constants.DUBLIN_BASE_URL,
+        ],
+        required=False,
+        default="api.segment.io/",
+    )
+
     class Meta:
         model = SegmentConfiguration
-        fields = ("id", "api_key")
+        fields = ("id", "api_key", "base_url")

api/tests/unit/integrations/segment/test_unit_segment.py

@@ -10,13 +10,15 @@
 def test_segment_initialized_correctly():
     # Given
     api_key = "123key"
-    config = SegmentConfiguration(api_key=api_key)
+    base_url = "api.segment.io/"
+    config = SegmentConfiguration(api_key=api_key, base_url=base_url)
 
-    # When initialized
+    # When
     segment_wrapper = SegmentWrapper(config)
 
     # Then
     assert segment_wrapper.analytics.write_key == api_key
+    assert segment_wrapper.analytics.host == base_url
 
 
 @pytest.mark.django_db

api/tests/unit/integrations/segment/test_unit_segment_views.py

@@ -5,6 +5,7 @@
 from rest_framework.test import APIClient
 
 from environments.models import Environment
+from integrations.segment.constants import DEFAULT_BASE_URL
 from integrations.segment.models import SegmentConfiguration
 
 
@@ -13,6 +14,7 @@ def test_should_create_segment_config_when_post(
     admin_client: APIClient,
 ) -> None:
     # Given
+
     data = {"api_key": "abc-123"}
     url = reverse(
         "api-v1:environments:integrations-segment-list",
@@ -30,7 +32,9 @@ def test_should_create_segment_config_when_post(
     assert response.status_code == status.HTTP_201_CREATED
     qs = SegmentConfiguration.objects.filter(environment=environment)
     assert qs.count() == 1
-    assert response.data["id"] == qs.first().id
+    segment_configuration = qs.first()
+    assert response.data["id"] == segment_configuration.id
+    assert segment_configuration.base_url == DEFAULT_BASE_URL
 
 
 def test_should_return_400_when_duplicate_segment_config_is_posted(
@@ -105,7 +109,7 @@ def test_should_return_segment_config_list_when_requested(
     response = admin_client.get(url)
 
     # Then
-    expected_response = {"api_key": config.api_key, "id": config.id}
+    expected_response = {"api_key": config.api_key, "id": config.id, "base_url": None}
     assert response.status_code == status.HTTP_200_OK
     assert response.data == [expected_response]