Ensuring we don't override existing sessions

Author: bsipocz

astroquery/query.py

@@ -185,11 +185,19 @@ class BaseVOQuery:
     """
     def __init__(self):
         super().__init__()
-        self._session = requests.Session()
-        self._session.headers['User-Agent'] = (
-            f"astroquery/{version.version} pyVO/{pyvo.__version__} Python/{platform.python_version()} "
-            f"({platform.system()}) "
-            f"{self._session.headers['User-Agent']}")
+        if not hasattr(self, '_session'):
+            # We don't want to override another, e.g. already authenticated session from another baseclass
+            self._session = requests.Session()
+
+        user_agents = self._session.headers['User-Agent'].split()
+        if 'astroquery' in user_agents[0]:
+            if 'pyVO' not in user_agents[1]:
+                user_agents[0] = f"astroquery/{version.version} pyVO/{pyvo.__version__}"
+        else:
+            user_agents = [f"astroquery/{version.version} pyVO/{pyvo.__version__} "
+                           f"Python/{platform.python_version()} ({platform.system()})"] + user_agents
+
+        self._session.headers['User-Agent'] = " ".join(user_agents)
 
         self.name = self.__class__.__name__.split("Class")[0]
 

astroquery/tests/test_query.py

@@ -0,0 +1,39 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+from astroquery.query import BaseQuery, BaseVOQuery
+
+
+class with_VO(BaseVOQuery, BaseQuery):
+    pass
+
+
+class without_VO(BaseQuery):
+    pass
+
+
+class only_VO(BaseVOQuery):
+    pass
+
+
+def test_session_VO_header():
+    test_instance = with_VO()
+    user_agent = test_instance._session.headers['User-Agent']
+    assert 'astroquery' in user_agent
+    assert 'pyVO' in user_agent
+    assert user_agent.count('astroquery') == 1
+
+
+def test_session_nonVO_header():
+    test_instance = without_VO()
+    user_agent = test_instance._session.headers['User-Agent']
+    assert 'astroquery' in user_agent
+    assert 'pyVO' not in user_agent
+    assert user_agent.count('astroquery') == 1
+
+
+def test_session_hooks():
+    # Test that we don't override the session in the BaseVOQuery
+    test_instance = with_VO()
+    assert len(test_instance._session.hooks['response']) > 0
+
+    test_VO_instance = only_VO()
+    assert len(test_VO_instance._session.hooks['response']) == 0