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

mage_ai/tests/data_preparation/models/global_hooks/test_utils.py

@@ -0,0 +1,33 @@
+import uuid
+
+from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.data_preparation.models.global_hooks.constants import (
+    DISABLED_RESOURCE_TYPES,
+    RESTRICTED_RESOURCE_TYPES,
+    VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL,
+)
+from mage_ai.data_preparation.models.global_hooks.utils import extract_valid_data
+from mage_ai.shared.hash import ignore_keys
+from mage_ai.tests.base_test import AsyncDBTestCase
+
+
+class GlobalHookUtilsTest(AsyncDBTestCase):
+    def test_extract_valid_data(self):
+        data = dict(
+            mage=uuid.uuid4().hex,
+        )
+
+        for key in VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL:
+            data[key] = uuid.uuid4().hex
+
+        types = DISABLED_RESOURCE_TYPES + RESTRICTED_RESOURCE_TYPES
+        for entity_name in EntityName:
+            if entity_name in [EntityName.ALL, EntityName.ALL_EXCEPT_RESERVED]:
+                continue
+
+            extracted_data = extract_valid_data(entity_name, data)
+
+            if entity_name in types:
+                self.assertEqual(extracted_data, dict(hook=data['hook']))
+            else:
+                self.assertEqual(extracted_data, ignore_keys(data, ['mage']))

mage_ai/tests/data_preparation/models/test_pipeline.py

@@ -121,6 +121,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[
@@ -266,6 +267,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[],
@@ -373,6 +375,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[],
@@ -537,6 +540,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[],
@@ -704,6 +708,7 @@ class PipelineTest(DBTestCase):
                     ],
                     callbacks=[],
                     conditionals=[],
+                    settings=dict(triggers=None),
                     widgets=[],
                 ),
             )

mage_ai/tests/data_preparation/test_repo_manager.py

@@ -66,6 +66,17 @@ class RepoManagerTest(DBTestCase):
         )
         shutil.rmtree(test.variables_dir)
 
+    def test_pipelines(self):
+        test = RepoConfig(repo_path=os.path.join(self.repo_path, 'non_existing_path'))
+        test.pipelines = dict(
+            settings=dict(
+                triggers=dict(
+                    save_in_code_automatically=True,
+                ),
+            ),
+        )
+        self.assertTrue(test.pipelines.settings.triggers.save_in_code_automatically)
+
     def test_variables_dir_expanduser(self):
         dir_name = uuid.uuid4().hex
         metadata_dict = dict(

mage_ai/tests/factory.py

@@ -6,12 +6,12 @@ from faker import Faker
 
 from mage_ai.authentication.passwords import create_bcrypt_hash, generate_salt
 from mage_ai.data_preparation.models.block import Block
-from mage_ai.data_preparation.models.constants import PipelineType
+from mage_ai.data_preparation.models.constants import BlockType, PipelineType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import ScheduleType
 from mage_ai.orchestration.db.models.oauth import User
 from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
-from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.hash import extract, ignore_keys, merge_dict
 
 faker = Faker()
 
@@ -24,7 +24,7 @@ def create_pipeline(name: str, repo_path: str):
     return pipeline
 
 
-def create_pipeline_with_blocks(name: str, repo_path: str):
+def create_pipeline_with_blocks(name: str, repo_path: str, pipeline_type: PipelineType = None):
     """
     Creates a pipeline with blocks for data processing and transformation.
 
@@ -38,6 +38,7 @@ def create_pipeline_with_blocks(name: str, repo_path: str):
     pipeline = Pipeline.create(
         name,
         repo_path=repo_path,
+        pipeline_type=pipeline_type,
     )
     block1 = Block.create('block1', 'data_loader', repo_path, language='python')
     block2 = Block.create('block2', 'transformer', repo_path, language='python')
@@ -188,3 +189,64 @@ def create_user(
         user.save()
 
     return user
+
+
+async def build_pipeline_with_blocks_and_content(
+    test_case,
+    block_settings: Dict = None,
+    name: str = None,
+    pipeline_type: PipelineType = None,
+) -> Pipeline:
+    repo_path = test_case.repo_path
+
+    pipeline = Pipeline.create(
+        name or test_case.faker.unique.name(),
+        repo_path=repo_path,
+        pipeline_type=pipeline_type or PipelineType.PYTHON,
+    )
+
+    blocks = []
+    for idx, block_dict in enumerate([
+        {},
+        {},
+        {},
+        {},
+    ]):
+        block = Block.create(
+            f'block{idx}_{test_case.faker.unique.name()}',
+            **merge_dict(
+                merge_dict(dict(
+                    block_type=BlockType.DATA_LOADER,
+                    repo_path=repo_path,
+                ), block_dict or {}),
+                ignore_keys(
+                    block_settings[idx] if block_settings and idx in block_settings else {},
+                    [
+                        'content',
+                    ],
+                ),
+            ),
+        )
+
+        pipeline.add_block(block)
+        blocks.append(block)
+
+    pipeline.save()
+    await pipeline.update(dict(
+        blocks=[merge_dict(
+            block.to_dict(include_content=True),
+            extract(
+                block_settings[idx] if block_settings and idx in block_settings else {},
+                [
+                    'content',
+                ],
+            ),
+        ) for idx, block in enumerate(blocks)],
+    ), update_content=True)
+
+    if hasattr(test_case, 'pipelines_created_for_testing'):
+        if not test_case.pipelines_created_for_testing:
+            test_case.pipelines_created_for_testing = []
+        test_case.pipelines_created_for_testing.append(pipeline)
+
+    return pipeline, blocks

mage_ai/tests/orchestration/queue/test_process_queue.py

@@ -1,8 +1,14 @@
+from unittest.mock import patch
+
 from mage_ai.orchestration.queue.config import QueueConfig
 from mage_ai.orchestration.queue.process_queue import JobStatus, ProcessQueue
 from mage_ai.tests.base_test import TestCase
 
 
+def run_block():
+    print('test run block')
+
+
 class ProcessQueueTests(TestCase):
     def setUp(self):
         queue_config = QueueConfig.load(config=dict(concurrency=100))
@@ -17,17 +23,24 @@ class ProcessQueueTests(TestCase):
         self.queue.job_dict['block_run_3'] = JobStatus.COMPLETED
         self.queue.job_dict['block_run_4'] = JobStatus.CANCELLED
         self.queue.clean_up_jobs()
-        self.assertEqual(self.queue.job_dict['block_run_1'], JobStatus.QUEUED)
+        # queue is empty, thus 'block_run_1' is not in queue
+        self.assertFalse('block_run_1' in self.queue.job_dict)
         self.assertEqual(self.queue.job_dict['block_run_2'], 100)
         self.assertFalse('block_run_3' in self.queue.job_dict)
         self.assertFalse('block_run_4' in self.queue.job_dict)
 
-    def test_has_job(self):
+    @patch('mage_ai.orchestration.queue.process_queue.poll_job_and_execute')
+    def test_has_job(self, mock_poll_job_and_execute):
         self.queue.job_dict['block_run_1'] = JobStatus.QUEUED
         self.queue.job_dict['block_run_2'] = 100
         self.queue.job_dict['block_run_3'] = JobStatus.COMPLETED
         self.queue.job_dict['block_run_4'] = JobStatus.CANCELLED
+        # Queue is empty, thus return False
+        self.assertFalse(self.queue.has_job('block_run_1'))
+        # After enqueueing the job, the has_job method returns True
+        self.queue.enqueue('block_run_1', run_block)
         self.assertTrue(self.queue.has_job('block_run_1'))
+
         self.assertTrue(self.queue.has_job('block_run_2'))
         self.assertFalse(self.queue.has_job('block_run_3'))
         self.assertFalse(self.queue.has_job('block_run_4'))

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)