mage-ai 0.9.69__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/tests/data_preparation/test_variable_manager.py

@@ -5,7 +5,7 @@ from pandas.testing import assert_frame_equal
 
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.pipeline import Pipeline
-from mage_ai.data_preparation.models.variable import VariableType
+from mage_ai.data_preparation.models.variables.constants import VariableType
 from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.data_preparation.variable_manager import (
     VariableManager,
@@ -20,20 +20,23 @@ from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
 class VariableManagerTest(DBTestCase):
     def test_add_and_get_variable(self):
-        self.__create_pipeline('test pipeline 1')
+        self.__create_pipeline("test pipeline 1")
         variable_manager = VariableManager(
             variables_dir=get_repo_config(self.repo_path).variables_dir,
         )
-        data1 = {'k1': 'v1', 'k2': 'v2'}
-        data2 = pd.DataFrame([
-            ['test1', 1],
-            ['test2', 2],
-        ], columns=['col1', 'col2'])
+        data1 = {"k1": "v1", "k2": "v2"}
+        data2 = pd.DataFrame(
+            [
+                ["test1", 1],
+                ["test2", 2],
+            ],
+            columns=["col1", "col2"],
+        )
         data3 = dict(
             metadata=dict(
                 column_types=dict(
-                    col1='number',
-                    col2='text',
+                    col1="number",
+                    col2="text",
                 ),
             ),
             statistics=dict(
@@ -43,71 +46,73 @@ class VariableManagerTest(DBTestCase):
             insights=dict(),
             suggestions=[
                 dict(
-                    title='Remove outliers',
+                    title="Remove outliers",
                 )
-            ]
+            ],
         )
-        variable_manager.add_variable('test_pipeline_1', 'block1', 'var1', data1)
-        variable_manager.add_variable('test_pipeline_1', 'block2', 'var2', data2)
+        variable_manager.add_variable("test_pipeline_1", "block1", "var1", data1)
+        variable_manager.add_variable("test_pipeline_1", "block2", "var2", data2)
         variable_manager.add_variable(
-            'test_pipeline_1',
-            'block2',
-            'var2',
+            "test_pipeline_1",
+            "block2",
+            "var2",
             data3,
             variable_type=VariableType.DATAFRAME_ANALYSIS,
         )
         self.assertEqual(
-            variable_manager.get_variable('test_pipeline_1', 'block1', 'var1'),
+            variable_manager.get_variable("test_pipeline_1", "block1", "var1"),
             data1,
         )
         assert_frame_equal(
-            variable_manager.get_variable('test_pipeline_1', 'block2', 'var2'),
+            variable_manager.get_variable("test_pipeline_1", "block2", "var2"),
             data2,
         )
         self.assertEqual(
             variable_manager.get_variable(
-                'test_pipeline_1',
-                'block2',
-                'var2',
+                "test_pipeline_1",
+                "block2",
+                "var2",
                 variable_type=VariableType.DATAFRAME_ANALYSIS,
             ),
             data3,
         )
 
     def test_get_variables_by_pipeline(self):
-        self.__create_pipeline('test pipeline 2')
+        self.__create_pipeline("test pipeline 2")
         variable_manager = VariableManager(
             repo_path=self.repo_path,
             variables_dir=get_repo_config(self.repo_path).variables_dir,
         )
-        variable_manager.add_variable('test_pipeline_2', 'block1', 'var1', 1)
-        variable_manager.add_variable('test_pipeline_2', 'block1', 'var2', 2)
-        variable_manager.add_variable('test_pipeline_2', 'block2', 'var3', 3)
-        variable_manager.add_variable('test_pipeline_2', 'block2', 'var4', 4)
+        variable_manager.add_variable("test_pipeline_2", "block1", "var1", 1)
+        variable_manager.add_variable("test_pipeline_2", "block1", "var2", 2)
+        variable_manager.add_variable("test_pipeline_2", "block2", "var3", 3)
+        variable_manager.add_variable("test_pipeline_2", "block2", "var4", 4)
         self.assertEqual(
-            variable_manager.get_variables_by_pipeline('test_pipeline_2'),
-            dict(block1=['var1', 'var2'], block2=['var3', 'var4']),
+            variable_manager.get_variables_by_pipeline("test_pipeline_2"),
+            dict(block1=["var1", "var2"], block2=["var3", "var4"]),
         )
 
     def test_set_and_get_global_variable(self):
         set_repo_path(self.repo_path)
-        self.__create_pipeline('test pipeline 3')
-        set_global_variable('test_pipeline_3', 'var1', 1)
-        set_global_variable('test_pipeline_3', 'var2', 'test')
-        set_global_variable('test_pipeline_3', 'var3', [1, 2, 3])
-        set_global_variable('test_pipeline_3', 'var4', dict(k1='v1', k2='v2'))
-        self.assertEqual(get_global_variable('test_pipeline_3', 'var1'), 1)
-        self.assertEqual(get_global_variable('test_pipeline_3', 'var2'), 'test')
-        self.assertEqual(get_global_variable('test_pipeline_3', 'var3'), [1, 2, 3])
-        self.assertEqual(get_global_variable('test_pipeline_3', 'var4'), dict(k1='v1', k2='v2'))
+        self.__create_pipeline("test pipeline 3")
+        set_global_variable("test_pipeline_3", "var1", 1)
+        set_global_variable("test_pipeline_3", "var2", "test")
+        set_global_variable("test_pipeline_3", "var3", [1, 2, 3])
+        set_global_variable("test_pipeline_3", "var4", dict(k1="v1", k2="v2"))
+        self.assertEqual(get_global_variable("test_pipeline_3", "var1"), 1)
+        self.assertEqual(get_global_variable("test_pipeline_3", "var2"), "test")
+        self.assertEqual(get_global_variable("test_pipeline_3", "var3"), [1, 2, 3])
+        self.assertEqual(
+            get_global_variable("test_pipeline_3", "var4"), dict(k1="v1", k2="v2")
+        )
 
     def __create_pipeline(self, name):
         pipeline = Pipeline.create(
             name,
             repo_path=self.repo_path,
         )
-        block1 = Block.create('block1', 'data_loader', self.repo_path)
-        block2 = Block.create('block2', 'transformer', self.repo_path)
+        block1 = Block.create("block1", "data_loader", self.repo_path)
+        block2 = Block.create("block2", "transformer", self.repo_path)
         pipeline.add_block(block1)
         pipeline.add_block(block2)
         return pipeline
@@ -119,44 +124,46 @@ class VariableManagerTest(DBTestCase):
         )
         pipeline.variables = self.faker.unique.name()
         pipeline.save()
-        self.assertEqual(get_global_variables(None, pipeline=pipeline), pipeline.variables)
+        self.assertEqual(
+            get_global_variables(None, pipeline=pipeline), pipeline.variables
+        )
         self.assertEqual(get_global_variables(pipeline.uuid), pipeline.variables)
 
 
 class VariableManagerProjectPlatformTests(ProjectPlatformMixin):
     def test_get_global_variable(self):
         with patch(
-            'mage_ai.data_preparation.variable_manager.project_platform_activated',
+            "mage_ai.data_preparation.variable_manager.project_platform_activated",
             lambda: True,
         ):
             with patch(
-                'mage_ai.data_preparation.models.pipeline.project_platform_activated',
+                "mage_ai.data_preparation.models.pipeline.project_platform_activated",
                 lambda: True,
             ):
                 for settings in self.repo_paths.values():
                     pipeline = Pipeline.create(
                         self.faker.unique.name(),
-                        repo_path=settings['full_path'],
+                        repo_path=settings["full_path"],
                     )
                     value = self.faker.unique.name()
                     pipeline.variables = dict(mage=value)
                     pipeline.save()
 
-                    self.assertEqual(get_global_variable(pipeline.uuid, 'mage'), value)
+                    self.assertEqual(get_global_variable(pipeline.uuid, "mage"), value)
 
     def test_get_global_variables(self):
         with patch(
-            'mage_ai.data_preparation.variable_manager.project_platform_activated',
+            "mage_ai.data_preparation.variable_manager.project_platform_activated",
             lambda: True,
         ):
             with patch(
-                'mage_ai.data_preparation.models.pipeline.project_platform_activated',
+                "mage_ai.data_preparation.models.pipeline.project_platform_activated",
                 lambda: True,
             ):
                 for settings in self.repo_paths.values():
                     pipeline = Pipeline.create(
                         self.faker.unique.name(),
-                        repo_path=settings['full_path'],
+                        repo_path=settings["full_path"],
                     )
                     pipeline.variables = self.faker.unique.name()
                     pipeline.save()
@@ -165,4 +172,6 @@ class VariableManagerProjectPlatformTests(ProjectPlatformMixin):
                         get_global_variables(None, pipeline=pipeline),
                         pipeline.variables,
                     )
-                    self.assertEqual(get_global_variables(pipeline.uuid), pipeline.variables)
+                    self.assertEqual(
+                        get_global_variables(pipeline.uuid), pipeline.variables
+                    )

mage_ai/tests/factory.py

@@ -1,6 +1,7 @@
 import os
 from datetime import datetime, timedelta
-from typing import Dict, Union
+from random import randint
+from typing import Dict, Optional, Union
 
 from faker import Faker
 
@@ -20,11 +21,24 @@ from mage_ai.shared.hash import extract, ignore_keys, merge_dict
 faker = Faker()
 
 
-def create_pipeline(name: str, repo_path: str):
-    pipeline = Pipeline.create(
-        name,
-        repo_path=repo_path,
-    )
+def create_pipeline(name: str, repo_path: str, pipeline_type: Optional[PipelineType] = None):
+    tries = 0
+    pipeline = None
+    while pipeline is None and tries < 100:
+        try:
+            name_use = name
+            if tries >= 1:
+                name_use = f'{name_use} {tries}'
+
+            pipeline = Pipeline.create(
+                name_use,
+                repo_path=repo_path,
+                pipeline_type=pipeline_type,
+            )
+        except Exception as err:
+            print(f'[ERROR] create_pipeline_with_blocks: {err}.')
+        tries += 1
+
     return pipeline
 
 
@@ -44,22 +58,7 @@ def create_pipeline_with_blocks(
     Returns:
         Pipeline: The created pipeline with added blocks.
     """
-    tries = 0
-    pipeline = None
-    while pipeline is None and tries < 100:
-        try:
-            name_use = name
-            if tries >= 1:
-                name_use = f'{name_use} {tries}'
-
-            pipeline = Pipeline.create(
-                name_use,
-                repo_path=repo_path,
-                pipeline_type=pipeline_type,
-            )
-        except Exception as err:
-            print(f'[ERROR] create_pipeline_with_blocks: {err}.')
-        tries += 1
+    pipeline = create_pipeline(name, 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')
@@ -97,11 +96,11 @@ def create_integration_pipeline_with_blocks(name: str, repo_path: str):
     )
     with open(block1.file_path, 'w') as f:
         f.write(
-            '''config:
+            """config:
   host: host
   port: 22
 source: sftp
-'''
+"""
         )
     block2 = Block.create(
         'test integration transform',
@@ -117,7 +116,7 @@ source: sftp
     )
     with open(block3.file_path, 'w') as f:
         f.write(
-            '''config:
+            """config:
   database: postgres
   schema: public
   host: localhost
@@ -125,7 +124,7 @@ source: sftp
   username: postgres
   password: postgres
 destination: postgresql
-'''
+"""
         )
     pipeline.add_block(block1)
     pipeline.add_block(block2, upstream_block_uuids=['block1'])
@@ -153,8 +152,8 @@ def create_pipeline_with_dynamic_blocks(name: str, repo_path: str):
     return pipeline
 
 
-def create_pipeline_run(pipeline_uuid: str, **kwargs):
-    pipeline_run = PipelineRun.create(pipeline_uuid='test_pipeline', **kwargs)
+def create_pipeline_run(pipeline_uuid: Optional[str] = None, **kwargs):
+    pipeline_run = PipelineRun.create(pipeline_uuid=pipeline_uuid or 'test_pipeline', **kwargs)
     return pipeline_run
 
 
@@ -275,10 +274,13 @@ async def build_pipeline_with_blocks_and_content(
         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 {}),
+                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 {},
                     [
@@ -292,17 +294,23 @@ async def build_pipeline_with_blocks_and_content(
         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)
+    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:
@@ -310,3 +318,16 @@ async def build_pipeline_with_blocks_and_content(
         test_case.pipelines_created_for_testing.append(pipeline)
 
     return pipeline, blocks
+
+
+def create_block(pipeline, name: Optional[str] = None):
+    block_type = [BlockType.DATA_LOADER, BlockType.TRANSFORMER, BlockType.DATA_EXPORTER][
+        randint(0, 2)
+    ]
+
+    return Block.create(
+        name or faker.unique.name(),
+        block_type,
+        pipeline.repo_path,
+        language='python',
+    )

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

@@ -1553,12 +1553,12 @@ class PipelineRunTests(DBTestCase):
         execution_date_str = execution_date.strftime(format='%Y%m%dT%H%M%S')
         execution_date_str = execution_date.strftime(format='%Y%m%dT%H%M%S')
         expected_file_path1 = os.path.join(
-            get_repo_config(self.repo_path).variables_dir,
+            get_repo_config(repo_path=self.repo_path).variables_dir,
             'pipelines/test_pipeline/.logs',
             f'{pipeline_run.pipeline_schedule_id}/{execution_date_str}/pipeline.log',
         )
         expected_file_path2 = os.path.join(
-            get_repo_config(self.repo_path).variables_dir,
+            get_repo_config(repo_path=self.repo_path).variables_dir,
             'pipelines/test_pipeline/.logs',
             f'{pipeline_run.pipeline_schedule_id}/{execution_date_str}/scheduler.log',
         )
@@ -1728,7 +1728,7 @@ class BlockRunTests(DBTestCase):
         execution_date_str = execution_date.strftime(format='%Y%m%dT%H%M%S')
         for b in pipeline_run.block_runs:
             expected_file_path = os.path.join(
-                get_repo_config(self.repo_path).variables_dir,
+                get_repo_config(repo_path=self.repo_path).variables_dir,
                 'pipelines/test_pipeline/.logs',
                 f'{pipeline_run.pipeline_schedule_id}/{execution_date_str}/{b.block_uuid}.log',
             )