mage-ai 0.9.44__py3-none-any.whl → 0.9.46__py3-none-any.whl

mage_ai/tests/services/aws/ecs/test_config.py

@@ -0,0 +1,144 @@
+from unittest.mock import MagicMock, patch
+
+from mage_ai.services.aws.ecs.config import EcsConfig
+from mage_ai.tests.base_test import TestCase
+
+
+class TestEcsConfig(TestCase):
+
+    @patch('os.getenv', return_value='mock_metadata_uri')
+    @patch('requests.get')
+    @patch('boto3.client')
+    @patch('boto3.resource')
+    def test_load_extra_config(self, mock_resource, mock_client, mock_requests, mock_os_getenv):
+        # Mock responses for requests.get
+        mock_response_task = MagicMock()
+        mock_response_task.json.return_value = {'Cluster': 'mock_cluster', 'Family': 'mock_family'}
+        mock_response_container = MagicMock()
+        mock_response_container.json.return_value = {'Name': 'mock_container_name'}
+
+        # Configure the side effects for the requests.get mocks
+        mock_requests.side_effect = [mock_response_container, mock_response_task]
+
+        # Mock response for boto3.client.describe_tasks
+        mock_task_config = {
+            'tasks': [{
+                'attachments': [{'type': 'ElasticNetworkInterface', 'details': [
+                    {'name': 'subnetId', 'value': 'mock_subnet'},
+                    {'name': 'networkInterfaceId', 'value': 'mock_network_interface'}
+                ]}]
+            }]
+        }
+        mock_client.return_value.describe_tasks.return_value = mock_task_config
+
+        # Mock response for boto3.resource.NetworkInterface
+        mock_network_interface = MagicMock()
+        mock_network_interface.groups = [{'GroupId': 'mock_security_group'}]
+        mock_resource.return_value.NetworkInterface.return_value = mock_network_interface
+
+        # Create an instance of EcsConfig and call the method to test
+        ecs_config = EcsConfig.load_extra_config()
+
+        # Assertions
+        self.assertEqual(ecs_config, {
+            'cluster': 'mock_cluster',
+            'subnets': ['mock_subnet'],
+            'security_groups': ['mock_security_group'],
+            'task_definition': 'mock_family',
+            'container_name': 'mock_container_name'
+        })
+
+    def test_get_task_config(self):
+        # assign_public_ip: True, enable_execute_command: True
+        ecs_config1 = EcsConfig(
+            task_definition='mock_task_def',
+            container_name='mock_container',
+            cluster='mock_cluster',
+            security_groups=['mock_sg'],
+            subnets=['mock_subnet'],
+            assign_public_ip=True,
+            cpu=512,
+            enable_execute_command=True,
+            launch_type='FARGATE',
+            memory=1024,
+            network_configuration=None,
+            tags=['tag1', 'tag2'],
+            wait_timeout=600
+        )
+        task_config1 = ecs_config1.get_task_config(command='mock_command')
+        self.assertEqual(task_config1, {
+            'cluster': 'mock_cluster',
+            'count': 1,
+            'enableExecuteCommand': True,
+            'launchType': 'FARGATE',
+            'networkConfiguration': {
+                'awsvpcConfiguration': {
+                    'subnets': ['mock_subnet'],
+                    'assignPublicIp': 'ENABLED',
+                    'securityGroups': ['mock_sg'],
+                }
+            },
+            'platformVersion': 'LATEST',
+            'propagateTags': 'TASK_DEFINITION',
+            'tags': ['tag1', 'tag2'],
+            'taskDefinition': 'mock_task_def',
+            'overrides': {
+                'containerOverrides': [
+                    {
+                        'name': 'mock_container',
+                        'command': ['mock_command'],
+                        'cpu': 512,
+                        'memory': 1024,
+                    },
+                ],
+                'cpu': '512',
+                'memory': '1024',
+            },
+        })
+
+        # assign_public_ip: False, enable_execute_command: False
+        ecs_config2 = EcsConfig(
+            task_definition='mock_task_def',
+            container_name='mock_container',
+            cluster='mock_cluster',
+            security_groups=['mock_sg'],
+            subnets=['mock_subnet'],
+            assign_public_ip=False,
+            cpu=512,
+            enable_execute_command=False,
+            launch_type='FARGATE',
+            memory=1024,
+            network_configuration=None,
+            tags=['tag1', 'tag2'],
+            wait_timeout=600
+        )
+        task_config2 = ecs_config2.get_task_config(command='mock_command')
+        self.assertEqual(task_config2, {
+            'cluster': 'mock_cluster',
+            'count': 1,
+            'enableExecuteCommand': False,
+            'launchType': 'FARGATE',
+            'networkConfiguration': {
+                'awsvpcConfiguration': {
+                    'subnets': ['mock_subnet'],
+                    'assignPublicIp': 'DISABLED',
+                    'securityGroups': ['mock_sg'],
+                }
+            },
+            'platformVersion': 'LATEST',
+            'propagateTags': 'TASK_DEFINITION',
+            'tags': ['tag1', 'tag2'],
+            'taskDefinition': 'mock_task_def',
+            'overrides': {
+                'containerOverrides': [
+                    {
+                        'name': 'mock_container',
+                        'command': ['mock_command'],
+                        'cpu': 512,
+                        'memory': 1024,
+                    },
+                ],
+                'cpu': '512',
+                'memory': '1024',
+            },
+        })

mage_ai/tests/services/k8s/test_job_manager.py

@@ -2,7 +2,14 @@ import os
 from unittest.mock import MagicMock, patch
 
 from kubernetes import client
-from kubernetes.client import V1Container, V1EnvVar, V1Pod, V1PodSpec
+from kubernetes.client import (
+    V1Container,
+    V1EnvFromSource,
+    V1EnvVar,
+    V1Pod,
+    V1PodSpec,
+    V1SecretEnvSource,
+)
 
 from mage_ai.services.k8s.config import K8sExecutorConfig
 from mage_ai.services.k8s.job_manager import (
@@ -22,6 +29,13 @@ MOCK_POD_CONFIG = V1Pod(
                 V1EnvVar(name='VAR1', value='VALUE1'),
                 V1EnvVar(name='VAR2', value='VALUE2'),
             ],
+            env_from=[
+                V1EnvFromSource(
+                    config_map_ref=None,
+                    prefix=None,
+                    secret_ref=V1SecretEnvSource(name='mysecret', optional=None),
+                ),
+            ],
             volume_mounts=[],
             resources=[],
         )],
@@ -167,6 +181,13 @@ class JobManagerTests(TestCase):
         self.assertEqual(pod_config.containers[0].image_pull_policy, 'IfNotPresent')
         self.assertEqual(pod_config.containers[0].image, 'test_image')
         self.assertEqual(pod_config.containers[0].resources, mock_v1_resource_requirements)
+        self.assertEqual(pod_config.containers[0].env_from, [
+                V1EnvFromSource(
+                    config_map_ref=None,
+                    prefix=None,
+                    secret_ref=V1SecretEnvSource(name='mysecret', optional=None),
+                ),
+        ])
 
         mock_client.V1ResourceRequirements.assert_called_once_with(
             limits=dict(

mage_ai/tests/shared/mixins.py

@@ -0,0 +1,291 @@
+import json
+import os
+import uuid
+from typing import Callable, Dict, List
+
+from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.data_preparation.models.constants import PipelineType
+from mage_ai.data_preparation.models.global_hooks.constants import HookOutputKey
+from mage_ai.data_preparation.models.global_hooks.models import (
+    GlobalHooks,
+    Hook,
+    HookCondition,
+    HookOperation,
+    HookOutputBlock,
+    HookOutputSettings,
+    HookPredicate,
+    HookStage,
+    HookStrategy,
+)
+from mage_ai.shared.hash import merge_dict
+from mage_ai.tests.api.operations.test_base import BaseApiTestCase
+from mage_ai.tests.factory import build_pipeline_with_blocks_and_content
+
+
+def build_content(query: Dict) -> str:
+    return f"""
+if 'data_loader' not in globals():
+    from mage_ai.data_preparation.decorators import data_loader
+
+
+if 'transformer' not in globals():
+    from mage_ai.data_preparation.decorators import transformer
+
+
+if 'data_exporter' not in globals():
+    from mage_ai.data_preparation.decorators import data_exporter
+
+
+@data_loader
+def load_data(*args, **kwargs):
+    return {json.dumps(query)}
+
+
+@transformer
+def transform(*args, **kwargs):
+    return {json.dumps(query)}
+
+
+@data_exporter
+def export_data(*args, **kwargs):
+    return {json.dumps(query)}
+"""
+
+
+def build_hooks(
+    test_case,
+    pipeline: PipelineType,
+    global_hooks: GlobalHooks = None,
+    hook_settings: Dict = None,
+    matching_predicate_resource: Dict = None,
+    operation_type: HookOperation = None,
+    resource_type: EntityName = None,
+    test_predicate_after: bool = False,
+    test_predicate_before: bool = False,
+    predicates_match: List[HookPredicate] = None,
+    predicates_miss: List[HookPredicate] = None,
+    snapshot: bool = True,
+) -> List[Hook]:
+    if not global_hooks:
+        global_hooks = GlobalHooks.load_from_file()
+
+    if matching_predicate_resource is None:
+        matching_predicate_resource = dict(
+            name=pipeline.name,
+            type=pipeline.type,
+        )
+
+    if predicates_match is None:
+        predicates_match = [[
+            HookPredicate.load(resource=matching_predicate_resource),
+        ]]
+
+    if predicates_miss is None:
+        predicates_miss = [[
+            HookPredicate.load(resource=dict(
+                name=uuid.uuid4().hex,
+                type=pipeline.type,
+            )),
+            predicates_match[0][0],
+        ]]
+
+    if not operation_type:
+        operation_type = HookOperation.DETAIL
+
+    if not resource_type:
+        resource_type = EntityName.Pipeline
+
+    hooks = []
+    hooks_match = []
+    hooks_miss = []
+    for idx, hook_dict in enumerate([
+        dict(
+            stages=[HookStage.BEFORE],
+        ),
+        dict(
+            conditions=[HookCondition.FAILURE],
+            stages=[HookStage.BEFORE],
+        ),
+        dict(
+            conditions=[HookCondition.SUCCESS],
+            stages=[HookStage.AFTER],
+        ),
+        dict(
+            conditions=[HookCondition.SUCCESS],
+            stages=[HookStage.AFTER],
+        ),
+    ]):
+        hook_payload = hook_dict.copy()
+        if hook_settings and idx in hook_settings:
+            hook_setting = hook_settings[idx]
+            hook_payload.update(hook_setting)
+
+        hook = Hook.load(
+            operation_type=operation_type,
+            predicates=predicates_match + predicates_miss,
+            resource_type=resource_type,
+            strategies=[HookStrategy.RAISE],
+            uuid=f'hook{idx}_{test_case.faker.unique.name()}',
+            **hook_payload,
+        )
+
+        hook_miss = Hook.load(
+            operation_type=hook.operation_type,
+            resource_type=hook.resource_type,
+            **hook.to_dict(),
+        )
+        hook_miss.uuid = test_case.faker.unique.name()
+        hook_miss.predicates = predicates_miss
+
+        if snapshot:
+            hook.snapshot()
+            hook_miss.snapshot()
+
+        hooks.append(hook)
+        hooks.append(hook_miss)
+        hooks_match.append(hook)
+        hooks_miss.append(hook_miss)
+
+        global_hooks.add_hook(hook)
+        global_hooks.add_hook(hook_miss)
+
+    global_hooks.save()
+
+    return global_hooks, hooks, hooks_match, hooks_miss
+
+
+class GlobalHooksMixin(BaseApiTestCase):
+    def setUp(self):
+        self.pipelines_created_for_testing = []
+
+    async def setUpAsync(
+        self,
+        block_settings: Dict = None,
+        hook_settings: Callable = None,
+        matching_predicate_resource: Dict = None,
+        operation_type: HookOperation = None,
+        pipeline_type: PipelineType = None,
+        predicates_match: List[HookPredicate] = None,
+        predicates_miss: List[HookPredicate] = None,
+        resource_type: EntityName = None,
+        snapshot: bool = True,
+    ):
+        block_settings_init = {
+            0: dict(content=build_content({
+                'type[]': [PipelineType.STREAMING.value],
+            })),
+            1: dict(content=build_content({
+                'type[]': [PipelineType.INTEGRATION.value],
+            })),
+            2: dict(content=build_content(dict(powers=dict(fire=1)))),
+            3: dict(content=build_content(dict(level=2))),
+        }
+
+        if block_settings:
+            block_settings = merge_dict(block_settings_init, block_settings)
+        else:
+            block_settings = block_settings_init
+
+        pipeline1, blocks1 = await build_pipeline_with_blocks_and_content(
+            self,
+            block_settings=block_settings,
+            pipeline_type=pipeline_type or PipelineType.PYTHON,
+        )
+        pipeline2, blocks2 = await build_pipeline_with_blocks_and_content(
+            self,
+            block_settings=block_settings,
+            pipeline_type=pipeline_type or PipelineType.PYTHON,
+        )
+        block11, block12, block13, block14 = blocks1
+        block21, block22, block23, block24 = blocks2
+
+        hook_settings_use = {
+            0: dict(
+                outputs=[
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block11.uuid),
+                        key=HookOutputKey.QUERY,
+                    ),
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block12.uuid),
+                        key=HookOutputKey.QUERY,
+                    ),
+                ],
+                pipeline=dict(
+                    uuid=pipeline1.uuid,
+                ),
+            ),
+            1: dict(
+                pipeline=dict(
+                    uuid=pipeline1.uuid,
+                ),
+            ),
+            2: dict(
+                outputs=[
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block13.uuid),
+                        key=HookOutputKey.METADATA,
+                    ),
+                ],
+                pipeline=dict(
+                    uuid=pipeline1.uuid,
+                ),
+            ),
+            3: dict(
+                outputs=[
+                    HookOutputSettings.load(
+                        block=HookOutputBlock.load(uuid=block24.uuid),
+                        key=HookOutputKey.METADATA,
+                        keys=['water'],
+                    ),
+                ],
+                pipeline=dict(
+                    uuid=pipeline2.uuid,
+                ),
+            ),
+        }
+
+        if hook_settings:
+            hook_settings_use = merge_dict(hook_settings_use, hook_settings(dict(
+                blocks1=blocks1,
+                blocks2=blocks2,
+                pipeline1=pipeline1,
+                pipeline2=pipeline2,
+            )))
+
+        global_hooks, hooks, hooks_match, hooks_miss = build_hooks(
+            self,
+            hook_settings=hook_settings_use,
+            matching_predicate_resource=matching_predicate_resource,
+            operation_type=operation_type,
+            pipeline=pipeline1,
+            predicates_match=predicates_match,
+            predicates_miss=predicates_miss,
+            resource_type=resource_type,
+            snapshot=snapshot,
+        )
+
+        self.blocks1 = blocks1
+        self.blocks2 = blocks2
+        self.global_hooks = global_hooks
+        self.hooks = hooks
+        self.hooks_match = hooks_match
+        self.hooks_miss = hooks_miss
+        self.pipeline1 = pipeline1
+        self.pipeline2 = pipeline2
+
+    def tearDown(self):
+        file_path = GlobalHooks.file_path()
+        if os.path.exists(file_path):
+            os.remove(file_path)
+
+        self.pipeline1.delete()
+        self.pipeline2.delete()
+
+        super().tearDown()
+
+    def bootstrap(self):
+        pass
+
+    def cleanup(self):
+        pass

mage_ai/tests/shared/test_hash.py

@@ -1,4 +1,4 @@
-from mage_ai.shared.hash import get_json_value
+from mage_ai.shared.hash import get_json_value, set_value
 from mage_ai.tests.base_test import TestCase
 
 
@@ -20,3 +20,19 @@ class HashTests(TestCase):
             get_json_value('{"k1": "v1", "k2": {"k21": "v21", "k22": "v22"} }', 'k2.k22'),
             'v22'
         )
+
+    def test_get_value(self):
+        obj = dict(l1=dict(power=1))
+        value = 'mage'
+
+        self.assertEqual(
+            set_value(obj, ['l1', 'l2', 'l3'], value),
+            dict(
+                l1=dict(
+                    l2=dict(
+                        l3=value,
+                    ),
+                    power=1,
+                ),
+            ),
+        )

mage_ai/tests/streaming/sources/test_activemq.py

@@ -0,0 +1,32 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sources.activemq import ActiveMQSource
+from mage_ai.tests.base_test import TestCase
+
+
+class ActiveMQTests(TestCase):
+    def test_init(self):
+        with patch.object(ActiveMQSource, 'init_client') as mock_init_client:
+            ActiveMQSource(dict(
+                connection_host='localhost',
+                connection_port='61613',
+                queue_name='test_queue',
+                username='guest',
+                password='guest',
+            ))
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        with patch.object(ActiveMQSource, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                ActiveMQSource(dict(
+                    connection_port='61613',
+                    queue_name='test_queue',
+                    username='guest',
+                    password='guest',
+                ))
+            self.assertTrue(
+                '__init__() missing 1 required positional argument: \'connection_host\''
+                in str(context.exception),
+            )
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/tests/streaming/sources/test_nats_js.py

@@ -0,0 +1,32 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sources.nats_js import NATSSource
+from mage_ai.tests.base_test import TestCase
+
+
+class NATSTests(TestCase):
+    def test_init(self):
+        with patch.object(NATSSource, 'init_client') as mock_init_client:
+            NATSSource(
+                dict(
+                    connector_type='nats',
+                    server_url='nats://localhost:4222',
+                    subject='test_subject',
+                    consumer_name='test_consumer',
+                    stream_name='test_stream',
+                )
+            )
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        invalid_config = {
+            'subject': 'test_subject',
+            'consumer_name': 'test_consumer',
+        }
+
+        with patch.object(NATSSource, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                NATSSource(invalid_config)
+            error_msg = "missing 2 required positional arguments: 'server_url' and 'stream_name'"
+            self.assertIn(error_msg, str(context.exception))
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/tests/streaming/sources/test_source_factory.py

@@ -1,13 +1,27 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sources.activemq import ActiveMQSource
 from mage_ai.streaming.sources.azure_event_hub import AzureEventHubSource
 from mage_ai.streaming.sources.google_cloud_pubsub import GoogleCloudPubSubSource
 from mage_ai.streaming.sources.kafka import KafkaSource
+from mage_ai.streaming.sources.nats_js import NATSSource
 from mage_ai.streaming.sources.rabbitmq import RabbitMQSource
 from mage_ai.streaming.sources.source_factory import SourceFactory
 from mage_ai.tests.base_test import TestCase
-from unittest.mock import patch
 
 
 class SourceFactoryTests(TestCase):
+    def test_get_source_nats(self):
+        with patch.object(NATSSource,
+                          '__init__',
+                          return_value=None) as mock_init:
+            config = dict(
+                connector_type='nats',
+            )
+            source = SourceFactory.get_source(config)
+            self.assertIsInstance(source, NATSSource)
+            mock_init.assert_called_once_with(config)
+
     def test_get_source_kafka(self):
         with patch.object(KafkaSource,
                           '__init__',
@@ -52,6 +66,17 @@ class SourceFactoryTests(TestCase):
             self.assertIsInstance(source, RabbitMQSource)
             mock_init.assert_called_once_with(config)
 
+    def test_get_source_activemq(self):
+        with patch.object(ActiveMQSource,
+                          '__init__',
+                          return_value=None) as mock_init:
+            config = dict(
+                connector_type='activemq',
+            )
+            source = SourceFactory.get_source(config)
+            self.assertIsInstance(source, ActiveMQSource)
+            mock_init.assert_called_once_with(config)
+
     def test_get_source_other(self):
         s = "Consuming data from random is not supported " \
             "in streaming pipelines yet."