added NotImplementedError to abstract class, added unit tests

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/telemetry/telemetry_client.py

@@ -104,15 +104,15 @@ class BaseTelemetryClient(ABC):
 
     @abstractmethod
     def export_initial_telemetry_log(self, driver_connection_params, user_agent):
-        pass
+        raise NotImplementedError("Subclasses must implement export_initial_telemetry_log")
 
     @abstractmethod
     def export_failure_log(self, error_name, error_message):
-        pass
+        raise NotImplementedError("Subclasses must implement export_failure_log")
 
     @abstractmethod
     def close(self):
-        pass
+        raise NotImplementedError("Subclasses must implement close")
 
 
 # A single instance of the no-op client that can be reused

tests/unit/test_telemetry.py

@@ -9,17 +9,22 @@
     initialize_telemetry_client,
     get_telemetry_client,
     _remove_telemetry_client,
+    TelemetryHelper,
+    BaseTelemetryClient
 )
 from databricks.sql.telemetry.models.enums import (
     AuthMech,
     DatabricksClientType,
+    AuthFlow,
 )
 from databricks.sql.telemetry.models.event import (
     DriverConnectionParameters,
     HostDetails,
 )
 from databricks.sql.auth.authenticators import (
     AccessTokenAuthProvider,
+    DatabricksOAuthProvider,
+    ExternalAuthProvider,
 )
 
 
@@ -257,6 +262,72 @@ def test_close(self, telemetry_client_setup):
 
         client._flush.assert_called_once()
 
+    @patch("requests.post")
+    def test_telemetry_request_callback_success(self, mock_post, telemetry_client_setup):
+        """Test successful telemetry request callback."""
+        client = telemetry_client_setup["client"]
+        
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+       
+        mock_future = MagicMock()
+        mock_future.result.return_value = mock_response
+       
+        client._telemetry_request_callback(mock_future)
+       
+        mock_future.result.assert_called_once()
+
+    @patch("requests.post")
+    def test_telemetry_request_callback_failure(self, mock_post, telemetry_client_setup):
+        """Test telemetry request callback with failure"""
+        client = telemetry_client_setup["client"]
+        
+        # Test with non-200 status code
+        mock_response = MagicMock()
+        mock_response.status_code = 500
+        future = MagicMock()
+        future.result.return_value = mock_response
+        client._telemetry_request_callback(future)
+
+        # Test with exception
+        future = MagicMock()
+        future.result.side_effect = Exception("Test error")
+        client._telemetry_request_callback(future)
+
+    def test_telemetry_client_exception_handling(self, telemetry_client_setup):
+        """Test exception handling in telemetry client methods."""
+        client = telemetry_client_setup["client"]
+        
+        # Test export_initial_telemetry_log with exception
+        with patch.object(client, '_export_event', side_effect=Exception("Test error")):
+            # Should not raise exception
+            client.export_initial_telemetry_log(MagicMock(), "test-agent")
+        
+        # Test export_failure_log with exception
+        with patch.object(client, '_export_event', side_effect=Exception("Test error")):
+            # Should not raise exception
+            client.export_failure_log("TestError", "Test error message")
+        
+        # Test _send_telemetry with exception
+        with patch.object(client._executor, 'submit', side_effect=Exception("Test error")):
+            # Should not raise exception
+            client._send_telemetry([MagicMock()])
+
+    def test_send_telemetry_thread_pool_failure(self, telemetry_client_setup):
+        """Test handling of thread pool submission failure"""
+        client = telemetry_client_setup["client"]
+        client._executor.submit.side_effect = Exception("Thread pool error")      
+        event = MagicMock()
+        client._send_telemetry([event])
+
+    def test_base_telemetry_client_abstract_methods(self):
+        """Test that BaseTelemetryClient cannot be instantiated without implementing all abstract methods"""
+        class TestBaseClient(BaseTelemetryClient):
+            pass
+
+        with pytest.raises(TypeError):
+            TestBaseClient()  # Can't instantiate abstract class
+
 
 class TestTelemetrySystem:
     """Tests for the telemetry system functions."""
@@ -317,4 +388,87 @@ def test_close_telemetry_client(self, telemetry_system_reset):
         _remove_telemetry_client(session_id_hex)
 
         client = get_telemetry_client(session_id_hex)
-        assert client is NOOP_TELEMETRY_CLIENT
+        assert client is NOOP_TELEMETRY_CLIENT
+
+    @patch("databricks.sql.telemetry.telemetry_client._handle_unhandled_exception")
+    def test_global_exception_hook(self, mock_handle_exception, telemetry_system_reset):
+        """Test that global exception hook is installed and handles exceptions."""
+        from databricks.sql.telemetry.telemetry_client import _install_exception_hook, _handle_unhandled_exception
+            
+        _install_exception_hook()
+        
+        test_exception = ValueError("Test exception")
+        _handle_unhandled_exception(type(test_exception), test_exception, None)
+        
+        mock_handle_exception.assert_called_once_with(type(test_exception), test_exception, None)
+
+
+class TestTelemetryHelper:
+    """Tests for the TelemetryHelper class."""
+
+    def test_get_driver_system_configuration(self):
+        """Test getting driver system configuration."""
+        config = TelemetryHelper.get_driver_system_configuration()
+          
+        assert isinstance(config.driver_name, str)
+        assert isinstance(config.driver_version, str)
+        assert isinstance(config.runtime_name, str)
+        assert isinstance(config.runtime_vendor, str)
+        assert isinstance(config.runtime_version, str)
+        assert isinstance(config.os_name, str)
+        assert isinstance(config.os_version, str)
+        assert isinstance(config.os_arch, str)
+        assert isinstance(config.locale_name, str)
+        assert isinstance(config.char_set_encoding, str)
+        
+        assert config.driver_name == "Databricks SQL Python Connector"
+        assert "Python" in config.runtime_name
+        assert config.runtime_vendor in ["CPython", "PyPy", "Jython", "IronPython"]
+        assert config.os_name in ["Darwin", "Linux", "Windows"]
+        
+        # Verify caching behavior
+        config2 = TelemetryHelper.get_driver_system_configuration()
+        assert config is config2  # Should return same instance
+
+    def test_get_auth_mechanism(self):
+        """Test getting auth mechanism for different auth providers."""
+        # Test PAT auth
+        pat_auth = AccessTokenAuthProvider("test-token")
+        assert TelemetryHelper.get_auth_mechanism(pat_auth) == AuthMech.PAT
+        
+        # Test OAuth auth
+        oauth_auth = MagicMock(spec=DatabricksOAuthProvider)
+        assert TelemetryHelper.get_auth_mechanism(oauth_auth) == AuthMech.DATABRICKS_OAUTH
+        
+        # Test External auth
+        external_auth = MagicMock(spec=ExternalAuthProvider)
+        assert TelemetryHelper.get_auth_mechanism(external_auth) == AuthMech.EXTERNAL_AUTH
+        
+        # Test None auth provider
+        assert TelemetryHelper.get_auth_mechanism(None) is None
+        
+        # Test unknown auth provider
+        unknown_auth = MagicMock()
+        assert TelemetryHelper.get_auth_mechanism(unknown_auth) == AuthMech.CLIENT_CERT
+
+    def test_get_auth_flow(self):
+        """Test getting auth flow for different OAuth providers."""
+        # Test OAuth with existing tokens
+        oauth_with_tokens = MagicMock(spec=DatabricksOAuthProvider)
+        oauth_with_tokens._access_token = "test-access-token"
+        oauth_with_tokens._refresh_token = "test-refresh-token"
+        assert TelemetryHelper.get_auth_flow(oauth_with_tokens) == AuthFlow.TOKEN_PASSTHROUGH
+        
+        # Test OAuth with browser-based auth
+        oauth_with_browser = MagicMock(spec=DatabricksOAuthProvider)
+        oauth_with_browser._access_token = None
+        oauth_with_browser._refresh_token = None
+        oauth_with_browser.oauth_manager = MagicMock()
+        assert TelemetryHelper.get_auth_flow(oauth_with_browser) == AuthFlow.BROWSER_BASED_AUTHENTICATION
+        
+        # Test non-OAuth provider
+        pat_auth = AccessTokenAuthProvider("test-token")
+        assert TelemetryHelper.get_auth_flow(pat_auth) is None
+        
+        # Test None auth provider
+        assert TelemetryHelper.get_auth_flow(None) is None