mage-ai 0.9.56__py3-none-any.whl → 0.9.57__py3-none-any.whl

mage_ai/tests/data_preparation/models/test_project.py

@@ -1,5 +1,5 @@
 import os
-from unittest.mock import patch
+from unittest.mock import call, patch
 
 import yaml
 
@@ -80,9 +80,19 @@ class ProjectTest(ProjectPlatformMixin, AsyncDBTestCase):
                         self.assertEqual(project.version, VERSION)
                         self.assertFalse(project.root_project)
 
-                        mock_get_repo_config.assert_called_once_with(
-                            repo_path=os.path.join(base_repo_path(), 'mage_platform'),
-                            root_project=False,
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[0],
+                            call(
+                                repo_path=os.path.join(base_repo_path(), 'mage_platform'),
+                                root_project=False,
+                            ),
+                        )
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[1],
+                            call(
+                                repo_path=base_repo_path(),
+                                root_project=True,
+                            ),
                         )
 
                     with patch(
@@ -96,9 +106,19 @@ class ProjectTest(ProjectPlatformMixin, AsyncDBTestCase):
                         self.assertEqual(project.version, VERSION)
                         self.assertTrue(project.root_project)
 
-                        mock_get_repo_config.assert_called_once_with(
-                            repo_path=base_repo_path(),
-                            root_project=True,
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[0],
+                            call(
+                                repo_path=base_repo_path(),
+                                root_project=True,
+                            ),
+                        )
+                        self.assertEqual(
+                            mock_get_repo_config.mock_calls[1],
+                            call(
+                                repo_path=base_repo_path(),
+                                root_project=True,
+                            ),
                         )
 
     def test_repo_path_for_database_query(self):

mage_ai/tests/orchestration/db/models/test_schedules.py

@@ -1791,6 +1791,10 @@ class BlockRunTests(DBTestCase):
             self.assertEqual(b.logs.get('path'), expected_file_path)
 
 
+@patch(
+    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
+    lambda: True,
+)
 class PipelineScheduleProjectPlatformTests(ProjectPlatformMixin):
     def test_repo_query(self):
         with patch(
@@ -1831,7 +1835,9 @@ class PipelineScheduleProjectPlatformTests(ProjectPlatformMixin):
                 self.assertIn(ps.id, ids)
 
     def test_pipeline(self):
-        with patch('mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform') as mock:
+        with patch(
+            'mage_ai.orchestration.db.models.schedules_project_platform.get_pipeline_from_platform',
+        ) as mock:
             pipeline_uuid = self.faker.unique.name()
             repo_path = self.faker.unique.name()
 
@@ -1847,6 +1853,10 @@ class PipelineScheduleProjectPlatformTests(ProjectPlatformMixin):
             )
 
 
+@patch(
+    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
+    lambda: True,
+)
 class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
     def test_pipeline(self):
         pipeline_uuid = self.faker.unique.name()
@@ -1864,7 +1874,9 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
             create_block_runs=False,
         )
 
-        with patch('mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform') as mock:
+        with patch(
+            'mage_ai.orchestration.db.models.schedules_project_platform.get_pipeline_from_platform',
+        ) as mock:
             pipeline_run.pipeline
 
             mock.assert_called_once_with(
@@ -1926,11 +1938,12 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
                 return fake_logger_manager
 
         with patch(
-            'mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform_async',
+            'mage_ai.orchestration.db.models.schedules_project_platform.'
+            'get_pipeline_from_platform_async',
             __get_pipeline_from_platform_async,
         ):
             with patch(
-                'mage_ai.orchestration.db.models.schedules.LoggerManagerFactory',
+                'mage_ai.orchestration.db.models.schedules_project_platform.LoggerManagerFactory',
                 FakeLoggerManagerFactory,
             ):
                 await pipeline_run.logs_async()
@@ -1951,7 +1964,7 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
         )
 
         with patch(
-            'mage_ai.orchestration.db.models.schedules.get_global_variables',
+            'mage_ai.orchestration.db.models.schedules_project_platform.get_global_variables',
             wraps=lambda pipeline_uuid, pipeline: {},
         ) as mock:
             with patch(
@@ -1963,6 +1976,10 @@ class PipelineRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
                 self.assertEqual(mock.mock_calls[0][2]['pipeline'].uuid, pipeline_run.pipeline_uuid)
 
 
+@patch(
+    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
+    lambda: True,
+)
 class BlockRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
     async def test_logs_async(self):
         value = self.faker.unique.name()
@@ -2028,11 +2045,12 @@ class BlockRunProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
                 return fake_logger_manager
 
         with patch(
-            'mage_ai.orchestration.db.models.schedules.get_pipeline_from_platform_async',
+            'mage_ai.orchestration.db.models.schedules_project_platform.' +
+            'get_pipeline_from_platform_async',
             __get_pipeline_from_platform_async,
         ):
             with patch(
-                'mage_ai.orchestration.db.models.schedules.LoggerManagerFactory',
+                'mage_ai.orchestration.db.models.schedules_project_platform.LoggerManagerFactory',
                 FakeLoggerManagerFactory,
             ):
                 await block_run.logs_async()

mage_ai/tests/orchestration/test_pipeline_scheduler.py

@@ -24,7 +24,6 @@ from mage_ai.orchestration.notification.sender import NotificationSender
 from mage_ai.orchestration.pipeline_scheduler import (
     PipelineScheduler,
     check_sla,
-    run_block,
     schedule_all,
 )
 from mage_ai.shared.array import find
@@ -37,7 +36,6 @@ from mage_ai.tests.factory import (
     create_pipeline_with_blocks,
     create_pipeline_with_dynamic_blocks,
 )
-from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
 
 class PipelineSchedulerTests(DBTestCase):
@@ -111,8 +109,8 @@ class PipelineSchedulerTests(DBTestCase):
         for b in block_runs:
             self.assertEqual(b.status, BlockRun.BlockRunStatus.CANCELLED)
 
-    @patch('mage_ai.orchestration.pipeline_scheduler.run_block')
-    @patch('mage_ai.orchestration.pipeline_scheduler.job_manager')
+    @patch('mage_ai.orchestration.pipeline_scheduler_original.run_block')
+    @patch('mage_ai.orchestration.pipeline_scheduler_original.job_manager')
     def test_schedule(self, mock_job_manager, mock_run_pipeline):
         pipeline_run = create_pipeline_run_with_schedule(pipeline_uuid='test_pipeline')
         scheduler = PipelineScheduler(pipeline_run=pipeline_run)
@@ -423,8 +421,8 @@ class PipelineSchedulerTests(DBTestCase):
                 pipeline_run=pipeline_run,
             )
 
-    @patch('mage_ai.orchestration.pipeline_scheduler.run_pipeline')
-    @patch('mage_ai.orchestration.pipeline_scheduler.job_manager')
+    @patch('mage_ai.orchestration.pipeline_scheduler_original.run_pipeline')
+    @patch('mage_ai.orchestration.pipeline_scheduler_original.job_manager')
     def test_schedule_streaming(self, mock_job_manager, mock_run_pipeline):
         pipeline = create_pipeline_with_blocks(
             'test pipeline 2',
@@ -941,7 +939,7 @@ class PipelineSchedulerTests(DBTestCase):
         )
 
     @freeze_time('2023-05-01 01:20:33')
-    @patch('mage_ai.orchestration.pipeline_scheduler.job_manager')
+    @patch('mage_ai.orchestration.pipeline_scheduler_original.job_manager')
     def test_pipeline_run_timeout(self, mock_job_manager):
         mock_job_manager.add_job = MagicMock()
         pipeline = create_pipeline_with_blocks(
@@ -989,7 +987,7 @@ class PipelineSchedulerTests(DBTestCase):
         self.assertEqual(pipeline_run3.status, PipelineRun.PipelineRunStatus.RUNNING)
 
     @freeze_time('2023-05-01 01:20:33')
-    @patch('mage_ai.orchestration.pipeline_scheduler.job_manager')
+    @patch('mage_ai.orchestration.pipeline_scheduler_original.job_manager')
     def test_block_run_timeout(self, mock_job_manager):
         mock_job_manager.add_job = MagicMock()
         pipeline = create_pipeline_with_blocks(
@@ -1039,256 +1037,3 @@ class PipelineSchedulerTests(DBTestCase):
         PipelineScheduler(pipeline_run=pipeline_run2).schedule()
         self.assertEqual(block_run.status, BlockRun.BlockRunStatus.FAILED)
         self.assertEqual(block_run2.status, BlockRun.BlockRunStatus.RUNNING)
-
-
-class PipelineSchedulerProjectPlatformTests(ProjectPlatformMixin, DBTestCase):
-    def test_init(self):
-        with patch(
-            'mage_ai.orchestration.pipeline_scheduler.project_platform_activated',
-            lambda: True,
-        ):
-            with patch(
-                'mage_ai.data_preparation.models.pipeline.project_platform_activated',
-                lambda: True,
-            ):
-                with patch(
-                    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
-                    lambda: True,
-                ):
-                    for settings in self.repo_paths.values():
-                        pipeline = create_pipeline_with_blocks(
-                            self.faker.unique.name(),
-                            repo_path=settings['full_path'],
-                        )
-                        pipeline_schedule = PipelineSchedule.create(
-                            name=self.faker.unique.name(),
-                            pipeline_uuid=pipeline.uuid,
-                            repo_path=settings['full_path'],
-                            schedule_type=ScheduleType.TIME,
-                        )
-                        pipeline_run = PipelineRun.create(
-                            pipeline_schedule_id=pipeline_schedule.id,
-                            pipeline_uuid=pipeline.uuid,
-                        )
-                        scheduler = PipelineScheduler(pipeline_run=pipeline_run)
-                        self.assertEqual(scheduler.pipeline.uuid, pipeline.uuid)
-
-    def test_run_block(self):
-        with patch(
-            'mage_ai.orchestration.pipeline_scheduler.project_platform_activated',
-            lambda: True,
-        ):
-            with patch(
-                'mage_ai.orchestration.pipeline_scheduler.project_platform_activated',
-                lambda: True,
-            ):
-                with patch(
-                    'mage_ai.data_preparation.models.pipeline.project_platform_activated',
-                    lambda: True,
-                ):
-                    with patch(
-                        'mage_ai.orchestration.db.models.schedules.project_platform_activated',
-                        lambda: True,
-                    ):
-                        for settings in self.repo_paths.values():
-                            full_path = settings['full_path']
-                            pipeline = create_pipeline_with_blocks(
-                                self.faker.unique.name(),
-                                repo_path=full_path,
-                            )
-                            pipeline_schedule = PipelineSchedule.create(
-                                name=self.faker.unique.name(),
-                                pipeline_uuid=pipeline.uuid,
-                                repo_path=full_path,
-                                schedule_type=ScheduleType.TIME,
-                            )
-                            pipeline_run = PipelineRun.create(
-                                pipeline_schedule_id=pipeline_schedule.id,
-                                pipeline_uuid=pipeline.uuid,
-                                status=PipelineRun.PipelineRunStatus.RUNNING,
-                            )
-
-                            for block_run in pipeline_run.block_runs:
-                                block_run.update(
-                                    status=BlockRun.BlockRunStatus.RUNNING,
-                                    started_at=datetime.utcnow() - timedelta(seconds=601),
-                                )
-
-                                class FakeExecutor:
-                                    def execute(cls, **kwargs):
-                                        pass
-
-                                fake_executor = FakeExecutor()
-
-                                class FakeExecutorFactory:
-                                    @classmethod
-                                    def get_block_executor(
-                                        cls,
-                                        pipeline,
-                                        block_uuid,
-                                        execution_partition,
-                                        block_uuid_test=block_run.block_uuid,
-                                        fake_executor=fake_executor,
-                                        pipeline_run=pipeline_run,
-                                        pipeline_test=pipeline,
-                                    ):
-                                        self.assertEqual(block_uuid, block_uuid_test)
-                                        self.assertEqual(
-                                            execution_partition,
-                                            pipeline_run.execution_partition,
-                                        )
-                                        self.assertEqual(pipeline.uuid, pipeline_test.uuid)
-
-                                        return fake_executor
-
-                                class FakeRepoConfig:
-                                    def __init__(self):
-                                        self.retry_config = {}
-
-                                fake_repo_config = FakeRepoConfig()
-
-                                def __get_repo_config(
-                                    repo_path,
-                                    fake_repo_config=fake_repo_config,
-                                    full_path=full_path,
-                                ):
-                                    self.assertEqual(repo_path, full_path)
-                                    return fake_repo_config
-
-                                with patch(
-                                    'mage_ai.orchestration.pipeline_scheduler.ExecutorFactory',
-                                    FakeExecutorFactory,
-                                ):
-                                    with patch(
-                                        'mage_ai.orchestration.pipeline_scheduler.get_repo_config',
-                                        __get_repo_config,
-                                    ):
-                                        run_block(
-                                            pipeline_run_id=pipeline_run.id,
-                                            block_run_id=block_run.id,
-                                            variables={},
-                                            tags={},
-                                        )
-
-    def test_check_sla(self):
-        with patch(
-            'mage_ai.orchestration.pipeline_scheduler.project_platform_activated',
-            lambda: True,
-        ):
-            with patch(
-                'mage_ai.data_preparation.models.pipeline.project_platform_activated',
-                lambda: True,
-            ):
-                pipeline_schedules = []
-                pipeline_runs = []
-                for settings in self.repo_paths.values():
-                    full_path = settings['full_path']
-                    pipeline = create_pipeline_with_blocks(
-                        self.faker.unique.name(),
-                        repo_path=full_path,
-                    )
-                    pipeline_schedule = PipelineSchedule.create(
-                        name=self.faker.unique.name(),
-                        pipeline_uuid=pipeline.uuid,
-                        repo_path=full_path,
-                        schedule_type=ScheduleType.TIME,
-                        status=ScheduleStatus.ACTIVE,
-                    )
-                    pipeline_run = PipelineRun.create(
-                        pipeline_schedule_id=pipeline_schedule.id,
-                        pipeline_uuid=pipeline.uuid,
-                        status=PipelineRun.PipelineRunStatus.RUNNING,
-                    )
-
-                    pipeline_schedules.append(pipeline_schedule)
-                    pipeline_runs.append(pipeline_run)
-
-        PipelineRunMock = MagicMock()
-
-        with patch(
-            'mage_ai.orchestration.pipeline_scheduler.PipelineRun',
-            PipelineRunMock,
-        ):
-            with patch(
-                'mage_ai.data_preparation.models.pipeline.project_platform_activated',
-                lambda: True,
-            ):
-                with patch(
-                    'mage_ai.orchestration.db.models.schedules.project_platform_activated',
-                    lambda: True,
-                ):
-                    with patch.object(PipelineRunMock, 'in_progress_runs') as mock:
-                        check_sla()
-                        mock.assert_called_once_with(set([s.id for s in pipeline_schedules]))
-
-    @freeze_time('2023-11-11 12:30:00')
-    def test_schedule_all(self):
-        with patch(
-            'mage_ai.orchestration.pipeline_scheduler.project_platform_activated',
-            lambda: True,
-        ):
-            with patch(
-                'mage_ai.orchestration.pipeline_scheduler.project_platform_activated',
-                lambda: True,
-            ):
-                with patch(
-                    'mage_ai.data_preparation.models.pipeline.project_platform_activated',
-                    lambda: True,
-                ):
-                    with patch(
-                        'mage_ai.orchestration.db.models.schedules.project_platform_activated',
-                        lambda: True,
-                    ):
-                        PipelineRun.query.delete()
-
-                        pipeline_schedules = []
-                        for settings in self.repo_paths.values():
-                            full_path = settings['full_path']
-                            pipeline = create_pipeline_with_blocks(
-                                self.faker.unique.name(),
-                                repo_path=full_path,
-                            )
-                            pipeline_schedule = PipelineSchedule.create(
-                                name=self.faker.unique.name(),
-                                pipeline_uuid=pipeline.uuid,
-                                repo_path=full_path,
-                                schedule_interval='@hourly',
-                                schedule_type=ScheduleType.TIME,
-                                start_time=datetime(2023, 10, 10, 13, 13, 20),
-                                status=ScheduleStatus.ACTIVE,
-                            )
-                            pipeline_schedules.append(pipeline_schedule)
-
-                        count = PipelineRun.query.count()
-
-                        class PipelineSchedulerMock:
-                            def __init__(self, pipeline_run):
-                                self.pipeline_run = pipeline_run
-
-                            def schedule(self):
-                                pass
-
-                            def start(self):
-                                pass
-
-                        with patch(
-                            'mage_ai.orchestration.pipeline_scheduler.PipelineScheduler',
-                            PipelineSchedulerMock,
-                        ):
-                            with patch.object(PipelineSchedulerMock, 'schedule') as mock:
-                                with patch.object(PipelineSchedulerMock, 'start') as mock_start:
-                                    schedule_all()
-                                    count += 2
-                                    self.assertEqual(PipelineRun.query.count(), count)
-                                    self.assertEqual(mock.call_count, 0)
-                                    self.assertEqual(mock_start.call_count, 2)
-
-                            with patch.object(PipelineSchedulerMock, 'schedule') as mock:
-                                with patch.object(PipelineSchedulerMock, 'start') as mock_start:
-                                    for pr in PipelineRun.query.all():
-                                        pr.update(status=PipelineRun.PipelineRunStatus.RUNNING)
-
-                                    schedule_all()
-                                    self.assertEqual(PipelineRun.query.count(), count)
-                                    self.assertEqual(mock.call_count, 2)
-                                    self.assertEqual(mock_start.call_count, 0)