Allow proxy url in telemetry config to be optional to allow for local testing (#639)

Author: jo-basevi

payu/telemetry.py

@@ -26,11 +26,13 @@
 # Required telemetry configuration fields
 CONFIG_FIELDS = {
     "URL": "telemetry_url",
-    "PROXY_URL": "telemetry_proxy_url",
     "TOKEN": "telemetry_token",
     "SERVICE_NAME": "telemetry_service_name",
     "HOSTNAME": "hostname",
 }
+OPTIONAL_CONFIG_FIELDS = {
+    "PROXY_URL": "telemetry_proxy_url",
+}
 
 REQUEST_TIMEOUT = 10
 
@@ -185,12 +187,12 @@ def write_error_log(
 
 
 def post_telemetry_data(url: str,
-                        proxy_url: str,
                         token: str,
                         data: dict[str, Any],
                         service_name: str,
                         archive_path: Path,
                         job_file_path: Path,
+                        proxy_url: Optional[str] = None,
                         request_timeout: int = REQUEST_TIMEOUT,
                         ) -> None:
     """Posts telemetry data
@@ -199,8 +201,6 @@ def post_telemetry_data(url: str,
     ----------
     url: str
         Endpoint for the telemetry
-    proxy_url: str
-        Proxy URL for the telemetry record request
     token: str
         Header token for the telemetry request
     data: dict[str, Any]
@@ -213,6 +213,8 @@ def post_telemetry_data(url: str,
     job_file_path: Path
         Path to the job file for the run. This is used for writing error logs
         if the telemetry request fails
+    proxy_url: Optional[str]
+        Proxy URL for the telemetry record request
     request_timeout: int, default REQUEST_TIMEOUT
         Timeout while waiting for request
     """
@@ -228,13 +230,21 @@ def post_telemetry_data(url: str,
     }
 
     try:
-        response = requests.post(
-            url,
-            data=json.dumps(data),
-            headers=headers,
-            timeout=request_timeout,
-            proxies={"https": proxy_url, "http": proxy_url},
-        )
+        if proxy_url is None:
+            response = requests.post(
+                url,
+                data=json.dumps(data),
+                headers=headers,
+                timeout=request_timeout,
+            )
+        else:
+            response = requests.post(
+                url,
+                data=json.dumps(data),
+                headers=headers,
+                timeout=request_timeout,
+                proxies={"https": proxy_url, "http": proxy_url},
+            )
         if response.status_code >= 400:
             error_message = (
                 "Error posting telemetry request: "
@@ -278,12 +288,12 @@ def record_telemetry(run_info: dict[str, Any],
         target=post_telemetry_data,
         kwargs={
             "url": external_config[CONFIG_FIELDS["URL"]],
-            "proxy_url": external_config[CONFIG_FIELDS["PROXY_URL"]],
             "token": external_config[CONFIG_FIELDS["TOKEN"]],
             "data": run_info,
             "service_name": external_config[CONFIG_FIELDS["SERVICE_NAME"]],
             "archive_path": archive_path,
             "job_file_path": job_file_path,
+            "proxy_url": external_config.get(OPTIONAL_CONFIG_FIELDS["PROXY_URL"]),
         },
     )
     thread.start()

test/test_telemetry.py

@@ -172,7 +172,6 @@ def test_get_external_telemetry_config_no_file(
     "hostname",
     "telemetry_service_name",
     "telemetry_token",
-    "telemetry_proxy_url"
 ])
 def test_get_external_telemetry_config_missing_fields(
             tmp_path, setup_env, config_path, missing_field
@@ -196,6 +195,26 @@ def test_get_external_telemetry_config_missing_fields(
     check_invalid_get_external_config(tmp_path, expected_error)
 
 
+def test_get_external_telemetry_config_missing_proxy_url(
+            tmp_path, setup_env, config_path
+        ):
+    config_data = {
+        "telemetry_url": "some/server/url",
+        "hostname": "gadi",
+        "telemetry_service_name": "payu",
+        "telemetry_token": "some_token",
+    }
+    with open(config_path, 'w') as f:
+        json.dump(config_data, f)
+
+    # Assert no error is raised when proxy URL is missing
+    result = get_external_telemetry_config(
+        archive_path=tmp_path / "archive",
+        job_file_path=tmp_path / "job_file.json"
+    )
+    assert result == config_data
+
+
 def test_get_external_telemetry_config_invalid_json(
             tmp_path, setup_env, config_path
         ):
@@ -536,3 +555,46 @@ def start_side_effect():
     with open(TELEMETRY_1_0_0_SCHEMA_PATH, "r") as f:
         schema = json.load(f)
     jsonschema.validate(sent_data, schema)
+
+
+@pytest.mark.parametrize("proxy_url,expected_proxies", [
+    (None, None),
+    ("http://proxy.example.com:8080", {
+        "http": "http://proxy.example.com:8080",
+        "https": "http://proxy.example.com:8080"
+    })
+])
+def test_post_telemetry_data(tmp_path, proxy_url, expected_proxies):
+    """Test the post_telemetry_data function with and without proxy URL"""
+
+    with patch('requests.post') as mock_post:
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"status": "success"}
+        mock_post.return_value = mock_response
+
+        post_telemetry_data(
+            url="http://example.com/telemetry",
+            token="test-token",
+            data={"key": "value"},
+            service_name="test-service",
+            proxy_url=proxy_url,
+            archive_path=tmp_path / "archive",
+            job_file_path=tmp_path / "job_file.json"
+        )
+        assert mock_post.called
+        args, kwargs = mock_post.call_args
+        assert args == ("http://example.com/telemetry",)
+        assert kwargs.get('headers') == {
+            'Content-type': 'application/json',
+            'Authorization': 'Token test-token',
+        }
+        assert kwargs.get('timeout') == 10
+        assert kwargs.get('proxies') == expected_proxies
+        assert kwargs.get('data') == json.dumps(
+            {
+                "service": "test-service",
+                "version": "1.0.0",
+                "telemetry": {"key": "value"}
+            }
+        )