mage-ai 0.9.44__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/data_preparation/test_templates.py

@@ -240,6 +240,7 @@ bootstrap_server: "localhost:9092"
 topic: topic_name
 consumer_group: unique_consumer_group
 include_metadata: false
+api_version: 0.10.2
 
 # Uncomment the config below to use SSL config
 # security_protocol: "SSL"

mage_ai/tests/factory.py

@@ -1,3 +1,4 @@
+import os
 from datetime import datetime
 from typing import Dict, Union
 
@@ -5,11 +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 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()
 
@@ -22,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.
 
@@ -36,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')
@@ -48,6 +51,63 @@ def create_pipeline_with_blocks(name: str, repo_path: str):
     return pipeline
 
 
+def create_integration_pipeline_with_blocks(name: str, repo_path: str):
+    pipeline = Pipeline.create(
+        name,
+        pipeline_type=PipelineType.INTEGRATION,
+        repo_path=repo_path,
+    )
+    with open(
+        os.path.join(os.path.dirname(__file__), 'test_data_integration_catalog.json'),
+        'r',
+    ) as f:
+        config = f.read()
+        with open(pipeline.catalog_config_path, 'w') as f:
+            f.write(config)
+    block1 = Block.create(
+        'test integration source',
+        'data_loader',
+        repo_path,
+        language='yaml',
+    )
+    with open(block1.file_path, 'w') as f:
+        f.write(
+            '''config:
+  host: host
+  port: 22
+source: sftp
+'''
+        )
+    block2 = Block.create(
+        'test integration transform',
+        'transformer',
+        repo_path,
+        language='python',
+    )
+    block3 = Block.create(
+        'test integration destination',
+        'data_exporter',
+        repo_path,
+        language='yaml',
+    )
+    with open(block3.file_path, 'w') as f:
+        f.write(
+            '''config:
+  database: postgres
+  schema: public
+  host: localhost
+  port: 5432
+  username: postgres
+  password: postgres
+destination: postgresql
+'''
+        )
+    pipeline.add_block(block1)
+    pipeline.add_block(block2, upstream_block_uuids=['block1'])
+    pipeline.add_block(block3, upstream_block_uuids=['block2'])
+    return pipeline
+
+
 def create_pipeline_with_dynamic_blocks(name: str, repo_path: str):
     pipeline = Pipeline.create(
         name,
@@ -60,10 +120,8 @@ def create_pipeline_with_dynamic_blocks(name: str, repo_path: str):
         language='python',
         configuration=dict(dynamic=True),
     )
-    block2 = Block.create(
-        'block2', 'transformer', repo_path, language='python')
-    block3 = Block.create(
-        'block3', 'data_exporter', repo_path, language='python')
+    block2 = Block.create('block2', 'transformer', repo_path, language='python')
+    block3 = Block.create('block3', 'data_exporter', repo_path, language='python')
     pipeline.add_block(block1)
     pipeline.add_block(block2, upstream_block_uuids=['block1'])
     pipeline.add_block(block3, upstream_block_uuids=['block2'])
@@ -110,10 +168,13 @@ def create_user(
     password = password or faker.password()
     password_salt = generate_salt()
     password_hash = create_bcrypt_hash(password, password_salt)
-    payload = merge_dict(dict(
-        email=faker.email(),
-        username=faker.name(),
-    ), kwargs)
+    payload = merge_dict(
+        dict(
+            email=faker.email(),
+            username=faker.name(),
+        ),
+        kwargs,
+    )
 
     if as_dict:
         payload.update(password=password)
@@ -128,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'))