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

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

@@ -0,0 +1,38 @@
+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'],
+                    resource_id=data['resource_id'],
+                    resource_parent_id=data['resource_parent_id'],
+                    user=data['user'],
+                ))
+            else:
+                self.assertEqual(extracted_data, ignore_keys(data, ['mage']))

mage_ai/tests/data_preparation/models/test_block.py

@@ -13,6 +13,7 @@ from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.data_preparation.variable_manager import VariableManager
 from mage_ai.tests.base_test import DBTestCase
+from mage_ai.tests.factory import create_integration_pipeline_with_blocks
 
 
 class BlockTest(DBTestCase):
@@ -707,3 +708,28 @@ def test_output(output, *args) -> None:
             block3.upstream_block_uuids[1],
             'test_data_loader_1',
         )
+
+    def test_output_variables_for_integration_pipeline_blocks(self):
+        pipeline = create_integration_pipeline_with_blocks(
+            'test integration pipeline',
+            repo_path=self.repo_path,
+        )
+        block = pipeline.get_block('test_integration_source')
+
+        df = pd.DataFrame(
+            [
+                [1, 'abc@xyz.com', '32132'],
+                [2, 'abc2@xyz.com', '12345'],
+                [3, 'test', '1234'],
+                [4, 'abc@test.net', 'abcde'],
+                [5, 'abc12345@', '54321'],
+                [6, 'abcdef@123.com', '56789'],
+            ],
+            columns=['id', 'email', 'zip_code'],
+        )
+
+        block.store_variables({
+            'output_sample_data_stream1': df
+        })
+        output_variables = block.output_variables()
+        self.assertTrue('output_sample_data_stream1' in output_variables)

mage_ai/tests/data_preparation/models/test_pipeline.py

@@ -49,6 +49,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -121,6 +123,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[
@@ -211,6 +214,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -266,6 +271,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[],
@@ -301,6 +307,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -373,6 +381,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[],
@@ -414,6 +423,8 @@ class PipelineTest(DBTestCase):
             run_pipeline_in_one_process=False,
             spark_config=dict(),
             type='python',
+            remote_variables_dir=None,
+            variables_dir=self.repo_path,
             blocks=[
                 dict(
                     language='python',
@@ -537,6 +548,7 @@ class PipelineTest(DBTestCase):
             ],
             callbacks=[],
             conditionals=[],
+            settings=dict(triggers=None),
             created_at='2023-08-01 08:08:24+00:00',
             updated_at=None,
             widgets=[],
@@ -666,6 +678,8 @@ class PipelineTest(DBTestCase):
                     type='integration',
                     updated_at=None,
                     uuid='test_pipeline_9',
+                    remote_variables_dir=None,
+                    variables_dir=self.repo_path,
                     blocks=[
                         dict(
                             all_upstream_blocks_executed=True,
@@ -704,6 +718,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

@@ -1,17 +1,21 @@
 import os
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Dict, Union
 
 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.orchestration.db.models.schedules import (
+    Backfill,
+    PipelineRun,
+    PipelineSchedule,
+)
+from mage_ai.shared.hash import extract, ignore_keys, merge_dict
 
 faker = Faker()
 
@@ -24,7 +28,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 +42,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')
@@ -158,6 +163,40 @@ def create_pipeline_run_with_schedule(
     return pipeline_run
 
 
+def create_backfill(
+    pipeline_uuid: str,
+    end_datetime: datetime = None,
+    interval_type: Backfill.IntervalType = Backfill.IntervalType.DAY,
+    interval_units: int = 1,
+    pipeline_schedule_id: int = None,
+    start_datetime: datetime = None,
+    **kwargs,
+):
+    if start_datetime is None:
+        start_datetime = datetime.today() - timedelta(days=1)
+    if end_datetime is None:
+        end_datetime = datetime.today()
+    if pipeline_schedule_id is None:
+        pipeline_schedule = PipelineSchedule.create(
+            name=f'{pipeline_uuid}_trigger',
+            pipeline_uuid=pipeline_uuid,
+            schedule_type=ScheduleType.TIME,
+        )
+        pipeline_schedule_id = pipeline_schedule.id
+    backfill = Backfill.create(
+        end_datetime=end_datetime,
+        interval_type=interval_type,
+        interval_units=interval_units,
+        pipeline_uuid=pipeline_uuid,
+        pipeline_schedule_id=pipeline_schedule_id,
+        start_datetime=start_datetime,
+        status=Backfill.Status.INITIAL,
+        **kwargs,
+    )
+
+    return backfill
+
+
 def create_user(
     as_dict: bool = False,
     save: bool = True,
@@ -188,3 +227,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/orchestration/test_pipeline_scheduler.py

@@ -14,12 +14,14 @@ from mage_ai.data_preparation.models.triggers import (
 )
 from mage_ai.data_preparation.preferences import get_preferences
 from mage_ai.data_preparation.variable_manager import VariableManager
+from mage_ai.orchestration.backfills.service import start_backfill
 from mage_ai.orchestration.db.models.schedules import (
+    Backfill,
     BlockRun,
     PipelineRun,
     PipelineSchedule,
 )
-from mage_ai.orchestration.job_manager import JobType
+from mage_ai.orchestration.job_manager import JobManager, JobType
 from mage_ai.orchestration.notification.sender import NotificationSender
 from mage_ai.orchestration.pipeline_scheduler import (
     PipelineScheduler,
@@ -30,6 +32,7 @@ from mage_ai.shared.array import find
 from mage_ai.shared.hash import ignore_keys, merge_dict
 from mage_ai.tests.base_test import DBTestCase
 from mage_ai.tests.factory import (
+    create_backfill,
     create_integration_pipeline_with_blocks,
     create_pipeline_run_with_schedule,
     create_pipeline_with_blocks,
@@ -141,6 +144,84 @@ class PipelineSchedulerTests(DBTestCase):
                 pipeline_run=pipeline_run,
             )
 
+    @freeze_time('2023-11-11 12:30:00')
+    def test_backfill_status_with_completed_pipeline_run(self):
+        pipeline_schedule = PipelineSchedule.create(
+            name='test_pipeline_trigger',
+            pipeline_uuid='test_pipeline',
+        )
+        backfill = create_backfill(
+            'test_pipeline',
+            end_datetime=datetime.today(),
+            pipeline_schedule_id=pipeline_schedule.id,
+            start_datetime=datetime.today(),
+        )
+        pipeline_runs = start_backfill(backfill)
+        pipeline_run = pipeline_runs[0]
+        pipeline_run.update(status=PipelineRun.PipelineRunStatus.RUNNING)
+        scheduler = PipelineScheduler(pipeline_run=pipeline_run)
+        with patch.object(JobManager, 'add_job') as mock_add_br_job:
+            scheduler.schedule()
+            self.assertEqual(
+                backfill.status,
+                Backfill.Status.RUNNING,
+            )
+            mock_add_br_job.assert_called()
+
+            for b in pipeline_run.block_runs:
+                b.update(status=BlockRun.BlockRunStatus.COMPLETED)
+            scheduler.schedule()
+            self.assertEqual(
+                backfill.status,
+                Backfill.Status.COMPLETED,
+            )
+
+    @freeze_time('2023-11-11 12:30:00')
+    def test_backfill_status_with_failed_pipeline_run(self):
+        pipeline_schedule = PipelineSchedule.create(
+            name='test_pipeline_trigger',
+            pipeline_uuid='test_pipeline',
+        )
+        backfill = create_backfill(
+            'test_pipeline',
+            end_datetime=datetime.today(),
+            pipeline_schedule_id=pipeline_schedule.id,
+            start_datetime=datetime.today(),
+        )
+        pipeline_runs = start_backfill(backfill)
+        pipeline_run = pipeline_runs[0]
+        scheduler = PipelineScheduler(pipeline_run=pipeline_run)
+
+        for b in pipeline_run.block_runs:
+            b.update(status=BlockRun.BlockRunStatus.FAILED)
+        scheduler.schedule()
+        # Backfill with 1 pipeline run that fails
+        self.assertEqual(
+            backfill.status,
+            Backfill.Status.FAILED,
+        )
+
+        # Retry failed pipeline run that was associated with the existing backfill
+        retried_pipeline_run = create_pipeline_run_with_schedule(
+            'test_pipeline',
+            execution_date=pipeline_run.execution_date,
+            pipeline_schedule_id=pipeline_schedule.id,
+        )
+        scheduler2 = PipelineScheduler(pipeline_run=retried_pipeline_run)
+        backfill.update(status=Backfill.Status.INITIAL)
+        self.assertEqual(
+            backfill.status,
+            Backfill.Status.INITIAL,
+        )
+        # Retried pipeline run successfully completes, so backfill status updates to "completed"
+        for b in retried_pipeline_run.block_runs:
+            b.update(status=BlockRun.BlockRunStatus.COMPLETED)
+        scheduler2.schedule()
+        self.assertEqual(
+            backfill.status,
+            Backfill.Status.COMPLETED,
+        )
+
     @freeze_time('2023-10-11 12:13:14')
     def test_schedule_all_with_integration_pipeline(self):
         integration_pipeline_schedule = PipelineSchedule.create(

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,11 +29,26 @@ 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=[],
         )],
         restart_policy='Never',
         service_account_name='old_service_account_name',
+        tolerations=[
+            {
+                'effect': 'NoExecute',
+                'key': 'node.kubernetes.io/not-ready',
+                'operator': 'Exists',
+                'toleration_seconds': 300,
+            },
+        ],
         volumes=[],
     )
 )
@@ -167,6 +189,21 @@ 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),
+                ),
+        ])
+        self.assertEqual(pod_config.tolerations, [
+            {
+                'effect': 'NoExecute',
+                'key': 'node.kubernetes.io/not-ready',
+                'operator': 'Exists',
+                'toleration_seconds': 300,
+            },
+        ])
 
         mock_client.V1ResourceRequirements.assert_called_once_with(
             limits=dict(