mage-ai 0.9.74__py3-none-any.whl → 0.9.79__py3-none-any.whl

mage_ai/tests/streaming/sinks/test_kafka.py

@@ -0,0 +1,130 @@
+from unittest.mock import Mock, patch
+
+from mage_ai.streaming.sinks.kafka import KafkaSink
+from mage_ai.tests.base_test import TestCase
+
+
+class KafkaSinkTests(TestCase):
+    def test_init(self):
+        with patch.object(KafkaSink, 'init_client') as mock_init_client:
+            KafkaSink(dict(
+                connector_type='kafka',
+                bootstrap_server='test_server',
+                topic='test_topic',
+            ))
+            mock_init_client.assert_called_once()
+
+    def test_init_with_oauth_sasl_config(self):
+        with patch.object(KafkaSink, 'init_client') as mock_init_client:
+            sink = KafkaSink(dict(
+                connector_type='kafka',
+                bootstrap_server='test_server',
+                topic='test_topic',
+                security_protocol='SASL_SSL',
+                sasl_config=dict(
+                    mechanism='OAUTHBEARER',
+                    oauth_token_url='https://auth.example.com/token',
+                    oauth_client_id='test_client',
+                    oauth_client_secret='test_secret',
+                    oauth_scope='kafka',
+                ),
+            ))
+            mock_init_client.assert_called_once()
+            self.assertEqual(sink.config.security_protocol, 'SASL_SSL')
+            self.assertEqual(sink.config.sasl_config.mechanism, 'OAUTHBEARER')
+            self.assertEqual(
+                sink.config.sasl_config.oauth_token_url, 'https://auth.example.com/token')
+            self.assertEqual(sink.config.sasl_config.oauth_client_id, 'test_client')
+            self.assertEqual(sink.config.sasl_config.oauth_client_secret, 'test_secret')
+            self.assertEqual(sink.config.sasl_config.oauth_scope, 'kafka')
+            self.assertIsNone(sink.config.ssl_config)
+
+    def test_init_client_with_oauthbearer(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_token_url='https://auth.example.com/token',
+                oauth_client_id='test_client',
+                oauth_client_secret='test_secret',
+            ),
+        )
+
+        with patch('mage_ai.streaming.sinks.kafka.KafkaProducer') as mock_producer:
+            with patch('mage_ai.streaming.sources.kafka_oauth.requests.post') as mock_post:
+                # Mock the OAuth token response
+                mock_response = Mock()
+                mock_response.json.return_value = {
+                    'access_token': 'test_token',
+                    'expires_in': 3600,
+                }
+                mock_response.raise_for_status = Mock()
+                mock_post.return_value = mock_response
+                
+                KafkaSink(kafka_config)
+
+        # Verify KafkaProducer was called with correct parameters
+        self.assertEqual(mock_producer.call_count, 1)
+        args, kwargs = mock_producer.call_args
+        self.assertEqual(kwargs['bootstrap_servers'], 'test_server')
+        self.assertEqual(kwargs['security_protocol'], 'SASL_SSL')
+        self.assertEqual(kwargs['sasl_mechanism'], 'OAUTHBEARER')
+        self.assertIn('sasl_oauth_token_provider', kwargs)
+        self.assertIsNotNone(kwargs['sasl_oauth_token_provider'])
+
+    def test_init_client_with_oauthbearer_missing_token_url(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_client_id='test_client',
+                oauth_client_secret='test_secret',
+            ),
+        )
+
+        with self.assertRaises(Exception) as context:
+            KafkaSink(kafka_config)
+
+        self.assertIn('oauth_token_url is required', str(context.exception))
+
+    def test_init_client_with_oauthbearer_missing_client_id(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_token_url='https://auth.example.com/token',
+                oauth_client_secret='test_secret',
+            ),
+        )
+
+        with self.assertRaises(Exception) as context:
+            KafkaSink(kafka_config)
+
+        self.assertIn('oauth_client_id is required', str(context.exception))
+
+    def test_init_client_with_oauthbearer_missing_client_secret(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_token_url='https://auth.example.com/token',
+                oauth_client_id='test_client',
+            ),
+        )
+
+        with self.assertRaises(Exception) as context:
+            KafkaSink(kafka_config)
+
+        self.assertIn('oauth_client_secret is required', str(context.exception))

mage_ai/tests/streaming/sources/test_kafka.py

@@ -1,4 +1,4 @@
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 from mage_ai.streaming.sources.kafka import KafkaSource
 from mage_ai.tests.base_test import TestCase
@@ -33,7 +33,7 @@ class KafkaTests(TestCase):
             api_version='0.10.2',
             auto_offset_reset='latest',
             max_partition_fetch_bytes=1048576,
-            enable_auto_commit=True,
+            enable_auto_commit=False,
         )
 
     def test_init_client_with_topics(self):
@@ -56,7 +56,7 @@ class KafkaTests(TestCase):
             api_version='0.10.2',
             auto_offset_reset='latest',
             max_partition_fetch_bytes=1048576,
-            enable_auto_commit=True,
+            enable_auto_commit=False,
         )
 
     def test_init_client_with_missing_topic_or_topics(self):
@@ -93,6 +93,128 @@ class KafkaTests(TestCase):
             self.assertEqual(source.config.sasl_config.password, 'test_password')
             self.assertIsNone(source.config.ssl_config)
 
+    def test_init_with_oauth_sasl_config(self):
+        with patch.object(KafkaSource, 'init_client') as mock_init_client:
+            source = KafkaSource(dict(
+                connector_type='kafka',
+                bootstrap_server='test_server',
+                consumer_group='test_group',
+                topic='test_topic',
+                security_protocol='SASL_SSL',
+                sasl_config=dict(
+                    mechanism='OAUTHBEARER',
+                    oauth_token_url='https://auth.example.com/token',
+                    oauth_client_id='test_client',
+                    oauth_client_secret='test_secret',
+                    oauth_scope='kafka',
+                ),
+            ))
+            mock_init_client.assert_called_once()
+            self.assertEqual(source.config.security_protocol, 'SASL_SSL')
+            self.assertEqual(source.config.sasl_config.mechanism, 'OAUTHBEARER')
+            self.assertEqual(
+                source.config.sasl_config.oauth_token_url, 'https://auth.example.com/token')
+            self.assertEqual(source.config.sasl_config.oauth_client_id, 'test_client')
+            self.assertEqual(source.config.sasl_config.oauth_client_secret, 'test_secret')
+            self.assertEqual(source.config.sasl_config.oauth_scope, 'kafka')
+            self.assertIsNone(source.config.ssl_config)
+
+    def test_init_client_with_oauthbearer(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            consumer_group='test_group',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_token_url='https://auth.example.com/token',
+                oauth_client_id='test_client',
+                oauth_client_secret='test_secret',
+            ),
+        )
+
+        with patch('mage_ai.streaming.sources.kafka.KafkaConsumer') as mock_consumer:
+            with patch('mage_ai.streaming.sources.kafka_oauth.requests.post') as mock_post:
+                # Mock the OAuth token response
+                mock_response = Mock()
+                mock_response.json.return_value = {
+                    'access_token': 'test_token',
+                    'expires_in': 3600,
+                }
+                mock_response.raise_for_status = Mock()
+                mock_post.return_value = mock_response
+                
+                KafkaSource(kafka_config)
+
+        # Verify KafkaConsumer was called with correct parameters
+        self.assertEqual(mock_consumer.call_count, 1)
+        args, kwargs = mock_consumer.call_args
+        self.assertEqual(args, ('test_topic',))
+        self.assertEqual(kwargs['group_id'], 'test_group')
+        self.assertEqual(kwargs['bootstrap_servers'], 'test_server')
+        self.assertEqual(kwargs['security_protocol'], 'SASL_SSL')
+        self.assertEqual(kwargs['sasl_mechanism'], 'OAUTHBEARER')
+        self.assertIn('sasl_oauth_token_provider', kwargs)
+        self.assertIsNotNone(kwargs['sasl_oauth_token_provider'])
+
+    def test_init_client_with_oauthbearer_missing_token_url(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            consumer_group='test_group',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_client_id='test_client',
+                oauth_client_secret='test_secret',
+            ),
+        )
+
+        with self.assertRaises(Exception) as context:
+            KafkaSource(kafka_config)
+
+        self.assertIn('oauth_token_url is required', str(context.exception))
+
+    def test_init_client_with_oauthbearer_missing_client_id(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            consumer_group='test_group',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_token_url='https://auth.example.com/token',
+                oauth_client_secret='test_secret',
+            ),
+        )
+
+        with self.assertRaises(Exception) as context:
+            KafkaSource(kafka_config)
+
+        self.assertIn('oauth_client_id is required', str(context.exception))
+
+    def test_init_client_with_oauthbearer_missing_client_secret(self):
+        kafka_config = dict(
+            connector_type='kafka',
+            bootstrap_server='test_server',
+            consumer_group='test_group',
+            topic='test_topic',
+            security_protocol='SASL_SSL',
+            sasl_config=dict(
+                mechanism='OAUTHBEARER',
+                oauth_token_url='https://auth.example.com/token',
+                oauth_client_id='test_client',
+            ),
+        )
+
+        with self.assertRaises(Exception) as context:
+            KafkaSource(kafka_config)
+
+        self.assertIn('oauth_client_secret is required', str(context.exception))
+
     def test_init_with_serde_config(self):
         with patch.object(KafkaSource, 'init_client') as mock_init_client:
             source = KafkaSource(dict(

mage_ai/tests/streaming/sources/test_kafka_oauth.py

@@ -0,0 +1,208 @@
+from unittest.mock import Mock, patch
+
+import requests
+
+from mage_ai.streaming.sources.kafka_oauth import ClientCredentialsTokenProvider
+from mage_ai.tests.base_test import TestCase
+
+
+class ClientCredentialsTokenProviderTests(TestCase):
+    def test_init(self):
+        """Test token provider initialization"""
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+            scope='kafka',
+        )
+
+        self.assertEqual(provider.token_url, 'https://auth.example.com/token')
+        self.assertEqual(provider.client_id, 'test_client')
+        self.assertEqual(provider.client_secret, 'test_secret')
+        self.assertEqual(provider.scope, 'kafka')
+        self.assertIsNone(provider._access_token)
+        self.assertIsNone(provider._token_expiry_time)
+
+    @patch('mage_ai.streaming.sources.kafka_oauth.requests.post')
+    def test_token_fetch(self, mock_post):
+        """Test successful token fetch"""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            'access_token': 'test_token_123',
+            'expires_in': 3600,
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+        )
+
+        token = provider.token()
+
+        self.assertEqual(token, 'test_token_123')
+        mock_post.assert_called_once_with(
+            'https://auth.example.com/token',
+            data={
+                'grant_type': 'client_credentials',
+                'client_id': 'test_client',
+                'client_secret': 'test_secret',
+            },
+            headers={'Content-Type': 'application/x-www-form-urlencoded'},
+            timeout=10,
+        )
+
+    @patch('mage_ai.streaming.sources.kafka_oauth.requests.post')
+    def test_token_fetch_with_scope(self, mock_post):
+        """Test token fetch with scope parameter"""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            'access_token': 'test_token_123',
+            'expires_in': 3600,
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+            scope='kafka',
+        )
+
+        token = provider.token()
+
+        self.assertEqual(token, 'test_token_123')
+        mock_post.assert_called_once_with(
+            'https://auth.example.com/token',
+            data={
+                'grant_type': 'client_credentials',
+                'client_id': 'test_client',
+                'client_secret': 'test_secret',
+                'scope': 'kafka',
+            },
+            headers={'Content-Type': 'application/x-www-form-urlencoded'},
+            timeout=10,
+        )
+
+    @patch('mage_ai.streaming.sources.kafka_oauth.requests.post')
+    def test_token_reuse(self, mock_post):
+        """Test that token is reused when not expired"""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            'access_token': 'test_token_123',
+            'expires_in': 3600,
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+        )
+
+        # First call should fetch token
+        token1 = provider.token()
+        self.assertEqual(token1, 'test_token_123')
+        self.assertEqual(mock_post.call_count, 1)
+
+        # Second call should reuse token
+        token2 = provider.token()
+        self.assertEqual(token2, 'test_token_123')
+        self.assertEqual(mock_post.call_count, 1)
+
+    @patch('mage_ai.streaming.sources.kafka_oauth.requests.post')
+    @patch('mage_ai.streaming.sources.kafka_oauth.time.time')
+    def test_token_refresh_on_expiry(self, mock_time, mock_post):
+        """Test that token is refreshed when expired"""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            'access_token': 'test_token_123',
+            'expires_in': 120,
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        # Start time
+        mock_time.return_value = 1000.0
+
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+            token_expiry_buffer=60,
+        )
+
+        # First call should fetch token
+        token1 = provider.token()
+        self.assertEqual(token1, 'test_token_123')
+        self.assertEqual(mock_post.call_count, 1)
+
+        # Move time forward but not past expiry buffer
+        mock_time.return_value = 1030.0
+        token2 = provider.token()
+        self.assertEqual(token2, 'test_token_123')
+        self.assertEqual(mock_post.call_count, 1)
+
+        # Move time forward past expiry buffer
+        # Token expires at 1000 + 120 = 1120, buffer is 60, so refresh at 1060
+        mock_response.json.return_value = {
+            'access_token': 'test_token_456',
+            'expires_in': 120,
+        }
+        mock_time.return_value = 1070.0
+        token3 = provider.token()
+        self.assertEqual(token3, 'test_token_456')
+        self.assertEqual(mock_post.call_count, 2)
+
+    @patch('mage_ai.streaming.sources.kafka_oauth.requests.post')
+    def test_token_fetch_failure(self, mock_post):
+        """Test handling of token fetch failure"""
+        mock_post.side_effect = requests.exceptions.RequestException('Network error')
+
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+        )
+
+        with self.assertRaises(RuntimeError) as context:
+            provider.token()
+
+        self.assertIn('Failed to fetch OAuth token', str(context.exception))
+
+    @patch('mage_ai.streaming.sources.kafka_oauth.requests.post')
+    def test_token_invalid_response(self, mock_post):
+        """Test handling of invalid token response"""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            'invalid': 'response',
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+        )
+
+        with self.assertRaises(RuntimeError) as context:
+            provider.token()
+
+        self.assertIn('Invalid token response format', str(context.exception))
+        self.assertIn('missing', str(context.exception))
+
+    def test_extensions(self):
+        """Test that extensions returns empty dict"""
+        provider = ClientCredentialsTokenProvider(
+            token_url='https://auth.example.com/token',
+            client_id='test_client',
+            client_secret='test_secret',
+        )
+
+        extensions = provider.extensions()
+        self.assertEqual(extensions, {})

mage_ai/tests/streaming/sources/test_kafka_raw_value.py

@@ -0,0 +1,105 @@
+from unittest import TestCase
+from unittest.mock import MagicMock, patch
+import json
+
+from mage_ai.streaming.sources.kafka import KafkaSource
+from mage_ai.streaming.sources.shared import SerializationMethod
+
+
+class KafkaRawValueTests(TestCase):
+    @patch('mage_ai.streaming.sources.kafka.KafkaConsumer')
+    def test_kafka_raw_value_key_handling(self, mock_kafka):
+        # Create a mock message with non-UTF8 key
+        mock_message = MagicMock()
+        mock_message.key = b'\xff\x83\x00'  # Invalid UTF-8 bytes
+        mock_message.value = b'test_value'
+        mock_message.partition = 1
+        mock_message.offset = 100
+        mock_message.timestamp = 1234567890
+        mock_message.topic = 'test_topic'
+
+        # Configure source with RAW_VALUE
+        config = {
+            'bootstrap_server': 'dummy',
+            'consumer_group': 'test',
+            'topic': 'test_topic',
+            'include_metadata': True,
+            'serde_config': {
+                'serialization_method': SerializationMethod.RAW_VALUE
+            }
+        }
+        
+        source = KafkaSource(config)
+        
+        # Test message conversion
+        result = source._convert_message(mock_message)
+        
+        self.assertIsInstance(result, dict)
+        self.assertIn('metadata', result)
+        self.assertIn('key', result['metadata'])
+        # With RAW_VALUE, key should remain as bytes
+        self.assertEqual(result['metadata']['key'], b'\xff\x83\x00')
+        self.assertEqual(result['data'], b'test_value')
+
+    @patch('mage_ai.streaming.sources.kafka.KafkaConsumer')
+    def test_kafka_utf8_key_handling(self, mock_kafka):
+        # Create a mock message with UTF-8 key
+        mock_message = MagicMock()
+        mock_message.key = 'test_key'.encode('utf-8')
+        mock_message.value = json.dumps({'value': 'test_value'}).encode('utf-8')
+        mock_message.partition = 1
+        mock_message.offset = 100
+        mock_message.timestamp = 1234567890
+        mock_message.topic = 'test_topic'
+
+        # Configure source without RAW_VALUE
+        config = {
+            'bootstrap_server': 'dummy',
+            'consumer_group': 'test',
+            'topic': 'test_topic',
+            'include_metadata': True
+        }
+        
+        source = KafkaSource(config)
+        
+        # Test message conversion
+        result = source._convert_message(mock_message)
+        
+        self.assertIsInstance(result, dict)
+        self.assertIn('metadata', result)
+        self.assertIn('key', result['metadata'])
+        # Without RAW_VALUE, key should be decoded to string
+        self.assertEqual(result['metadata']['key'], 'test_key')
+        self.assertIsInstance(result['data'], dict)
+        self.assertEqual(result['data']['value'], 'test_value')
+
+    @patch('mage_ai.streaming.sources.kafka.KafkaConsumer')
+    def test_kafka_non_utf8_key_without_raw_value(self, mock_kafka):
+        # Create a mock message with non-UTF8 key but without RAW_VALUE
+        mock_message = MagicMock()
+        mock_message.key = b'\xff\x83\x00'  # Invalid UTF-8 bytes
+        mock_message.value = json.dumps({'value': 'test_value'}).encode('utf-8')
+        mock_message.partition = 1
+        mock_message.offset = 100
+        mock_message.timestamp = 1234567890
+        mock_message.topic = 'test_topic'
+
+        # Configure source without RAW_VALUE
+        config = {
+            'bootstrap_server': 'dummy',
+            'consumer_group': 'test',
+            'topic': 'test_topic',
+            'include_metadata': True
+        }
+        
+        source = KafkaSource(config)
+        
+        # Test message conversion - should fallback to raw bytes
+        result = source._convert_message(mock_message)
+        
+        self.assertIsInstance(result, dict)
+        self.assertIn('metadata', result)
+        self.assertIn('key', result['metadata'])
+        # Should fallback to raw bytes when decode fails
+        self.assertEqual(result['metadata']['key'], b'\xff\x83\x00')
+        self.assertIsInstance(result['data'], dict)