Merge pull request astropy#3526 from cds-astro/add-user-agent-in-basevoquery

feat: add extra_user_agents to BaseVOQuery

Author: bsipocz

CHANGES.rst

@@ -119,6 +119,9 @@ Infrastructure, Utility and Other Changes and Additions
 
 - Versions of numpy <1.22 are no longer supported. [#3504]
 
+- ``BaseVOQuery`` now accepts a ``extra_user_agents`` parameter to allow the addition
+  of user agents on top of astroquery's ones [#3526]
+
 
 utils.tap
 ^^^^^^^^^

astroquery/query.py

@@ -183,7 +183,14 @@ class BaseVOQuery:
     Use in modules that rely on PyVO, either on its own or in combination with ``BaseQuery`` (be mindful
     about resolution order of base classes!).
     """
-    def __init__(self):
+    def __init__(self, *, extra_user_agents=None):
+        """Create an instance of BaseVOQuery.
+
+        Parameters
+        ----------
+        extra_user_agents : str | list(str)
+            Extra user agents to be added to the ones already provided by astroquery.
+        """
         super().__init__()
         if not hasattr(self, '_session'):
             # We don't want to override another, e.g. already authenticated session from another baseclass
@@ -199,6 +206,12 @@ def __init__(self):
             user_agents = [f"astroquery/{version.version} pyVO/{pyvo.__version__} "
                            f"Python/{platform.python_version()} ({platform.system()})"] + user_agents
 
+        if extra_user_agents:
+            if isinstance(extra_user_agents, str):
+                user_agents += [extra_user_agents]
+            else:
+                user_agents += extra_user_agents
+
         self._session.headers['User-Agent'] = " ".join(user_agents)
 
         self.name = self.__class__.__name__.split("Class")[0]

astroquery/tests/test_query.py

@@ -332,13 +332,21 @@ def test_download_file_remote_large(self, base_query, tmp_path, head_safe):
 
 def test_session_VO_header():
     """Test that the session header includes both astroquery and pyVO."""
-    test_instance = with_VO()
+    test_instance = with_VO(extra_user_agents="test")
     user_agent = test_instance._session.headers['User-Agent']
     assert 'astroquery' in user_agent
     assert 'pyVO' in user_agent
+    assert user_agent.endswith("test")
     assert user_agent.count('astroquery') == 1
 
 
+def test_session_only_VO_header():
+    """Test that we can add extra user agents."""
+    test_instance = only_VO(extra_user_agents=["agent1", "agent2"])
+    user_agent = test_instance._session.headers['User-Agent']
+    assert all([string in user_agent for string in ["astroquery", "pyVO", "agent1", "agent2"]])
+
+
 def test_session_nonVO_header():
     """Test that the session header includes astroquery but not pyVO."""
     test_instance = without_VO()