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

mage_ai/tests/data_preparation/models/test_blocks_helper.py

@@ -0,0 +1,114 @@
+import inspect
+from typing import Callable, Optional
+
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.tabular.mocks import create_dataframe
+from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
+from mage_ai.tests.api.operations.test_base import BaseApiTestCase
+from mage_ai.tests.factory import create_pipeline
+
+
+def load_dynamic_block_child_data(*args, **kwargs):
+    arr = [i + 10 for i in range(0, 2)]
+    return [
+        arr,
+        [dict(block_uuid=f'child_{i}') for i in arr],
+    ]
+
+
+def load(*args, **kwargs):
+    buckets = 10
+
+    for i in DataGenerator(range(buckets)):
+        n_rows = (i + 1) * 100
+        yield create_dataframe(
+            n_rows=n_rows,
+            use_pandas=False,
+        )
+
+
+def transform_data(generator, *args, **kwargs):
+    for batch in generator:
+        yield batch
+
+
+def load_generators(*args, **kwargs):
+    def load_data(index):
+        return create_dataframe(n_rows=(1 + index) * 100, use_pandas=False)
+
+    def measure_data(**kwargs) -> int:
+        return 12
+
+    data_generator = DataGenerator(
+        load_data=load_data,
+        measure_data=lambda _: measure_data(**kwargs),
+    )
+
+    for data in data_generator:
+        yield data
+
+
+def load_dataframes_for_dynamic_children(*args, **kwargs):
+    from mage_ai.data.tabular.mocks import create_dataframe
+
+    return [
+        [
+            create_dataframe(n_rows=100, use_pandas=True),
+            create_dataframe(n_rows=100, use_pandas=True),
+            create_dataframe(n_rows=100, use_pandas=True),
+        ],
+    ]
+
+
+def load_dataframe(*args, **kwargs):
+    from mage_ai.data.tabular.mocks import create_dataframe
+
+    return create_dataframe(n_rows=1200, use_pandas=False)
+
+
+def process_generator(generator, **kwargs):
+    for batch in generator:
+        gen = batch.generator(batch_size=75)
+        for sub in gen:
+            df = sub.deserialize()
+            yield df
+
+
+def passthrough(data, **kwargs):
+    return data
+
+
+class BlockHelperTest(BaseApiTestCase):
+    def setUp(self):
+        super().setUp()
+        self.pipeline = create_pipeline(self.faker.unique.name(), self.repo_path)
+
+    def create_block(
+        self,
+        name: Optional[str] = None,
+        content: Optional[str] = None,
+        func: Optional[Callable] = None,
+        **kwargs,
+    ):
+        block = Block.create(
+            name or self.faker.unique.name(),
+            BlockType.TRANSFORMER,
+            self.pipeline.repo_path,
+            language=BlockLanguage.PYTHON,
+            **kwargs,
+        )
+
+        if func:
+            content = inspect.getsource(func)
+
+        if content:
+            block.update_content(
+                '\n'.join([
+                    'from mage_ai.data.models.generator import DataGenerator',
+                    'from mage_ai.data.tabular.mocks import create_dataframe',
+                    f'@{block.type}\n{content}',
+                ])
+            )
+
+        return block

mage_ai/tests/data_preparation/models/test_global_data_product.py

@@ -27,12 +27,17 @@ class GlobalDataProductTest(DBTestCase):
                 'test pipeline',
                 repo_path=self.repo_path,
             )
-            self.pipeline.add_block(Block('data_loader', 'data_loader', BlockType.DATA_LOADER))
-            self.pipeline.add_block(Block('transformer', 'transformer', BlockType.TRANSFORMER))
             self.pipeline.add_block(
-                Block('data_exporter', 'data_exporter', BlockType.DATA_EXPORTER))
+                Block('data_loader', 'data_loader', BlockType.DATA_LOADER)
+            )
+            self.pipeline.add_block(
+                Block('transformer', 'transformer', BlockType.TRANSFORMER)
+            )
+            self.pipeline.add_block(
+                Block('data_exporter', 'data_exporter', BlockType.DATA_EXPORTER)
+            )
         except Exception:
-            self.pipeline = Pipeline.get('test_pipeline')
+            self.pipeline = Pipeline.get('test_pipeline', repo_path=self.repo_path)
 
         self.global_data_product = GlobalDataProduct(
             object_type='pipeline',
@@ -54,6 +59,7 @@ class GlobalDataProductTest(DBTestCase):
                 week_of_month=8,
                 week_of_year=9,
             ),
+            repo_path=self.repo_path,
             settings=dict(
                 data_exporter={},
                 data_loader=dict(partitions=1),
@@ -61,13 +67,15 @@ class GlobalDataProductTest(DBTestCase):
             ),
             uuid='mage',
         )
-        self.pipeline_schedule = fetch_or_create_pipeline_schedule(self.global_data_product)
+        self.pipeline_schedule = fetch_or_create_pipeline_schedule(
+            self.global_data_product
+        )
 
         self.file_path = os.path.join(
             self.repo_path,
             'global_data_products.yaml',
         )
-        self.global_data_product.save(file_path=self.file_path)
+        self.global_data_product.save()
 
     def tearDown(self):
         super().tearDown()
@@ -80,7 +88,7 @@ class GlobalDataProductTest(DBTestCase):
         )
 
     def test_load_all(self):
-        arr = GlobalDataProduct.load_all(file_path=self.file_path)
+        arr = GlobalDataProduct.load_all(self.repo_path)
         self.assertEqual(len(arr), 1)
 
         gdp = arr[0]
@@ -95,7 +103,7 @@ class GlobalDataProductTest(DBTestCase):
             self.assertEqual(getattr(gdp, key), getattr(self.global_data_product, key))
 
     def test_get(self):
-        gdp = GlobalDataProduct.get('mage', file_path=self.file_path)
+        gdp = GlobalDataProduct.get('mage', self.repo_path)
 
         for key in [
             'object_type',
@@ -343,12 +351,16 @@ class GlobalDataProductTest(DBTestCase):
         )
 
         pipeline_run1 = PipelineRun.create(
-            execution_date=datetime(2023, 10, 11, 2, 13, 13).replace(tzinfo=timezone.utc),
+            execution_date=datetime(2023, 10, 11, 2, 13, 13).replace(
+                tzinfo=timezone.utc
+            ),
             pipeline_schedule_id=self.pipeline_schedule.id,
             pipeline_uuid=self.global_data_product.pipeline.uuid,
             status=PipelineRun.PipelineRunStatus.COMPLETED,
         )
-        pipeline_run1.execution_date = pipeline_run1.execution_date.replace(tzinfo=timezone.utc)
+        pipeline_run1.execution_date = pipeline_run1.execution_date.replace(
+            tzinfo=timezone.utc
+        )
 
         self.global_data_product.outdated_after = dict(seconds=60 * 60 * 10)
         self.assertEqual(
@@ -408,6 +420,8 @@ class GlobalDataProductTest(DBTestCase):
             outdated_after=self.global_data_product.outdated_after,
             outdated_starting_at=self.global_data_product.outdated_starting_at,
             settings=self.global_data_product.settings,
+            project=self.global_data_product.project_name,
+            repo_path=self.global_data_product.repo_path,
         )
 
         self.assertEqual(
@@ -422,34 +436,37 @@ class GlobalDataProductTest(DBTestCase):
         )
 
     def test_delete(self):
-        arr = GlobalDataProduct.load_all(file_path=self.file_path)
+        arr = GlobalDataProduct.load_all(self.repo_path)
         self.assertEqual(len(arr), 1)
         self.global_data_product.delete()
-        arr = GlobalDataProduct.load_all(file_path=self.file_path)
+        arr = GlobalDataProduct.load_all(self.repo_path)
         self.assertEqual(len(arr), 0)
 
     def test_save(self):
-        arr = GlobalDataProduct.load_all(file_path=self.file_path)
+        arr = GlobalDataProduct.load_all(self.repo_path)
         self.assertEqual(len(arr), 1)
 
         self.global_data_product.delete()
-        arr = GlobalDataProduct.load_all(file_path=self.file_path)
+        arr = GlobalDataProduct.load_all(self.repo_path)
         self.assertEqual(len(arr), 0)
 
         self.global_data_product.save()
-        arr = GlobalDataProduct.load_all(file_path=self.file_path)
+        arr = GlobalDataProduct.load_all(self.repo_path)
         self.assertEqual(len(arr), 1)
 
     def test_update(self):
-        self.global_data_product.update(dict(
-            object_type='test1',
-            object_uuid='test2',
-            outdated_after=dict(seconds=777),
-            outdated_starting_at=dict(day_of_month=40),
-            settings=dict(mage=dict(partitions=3)),
-        ))
-
-        arr = GlobalDataProduct.load_all(file_path=self.file_path)
+        self.global_data_product.update(
+            dict(
+                object_type='test1',
+                object_uuid='test2',
+                outdated_after=dict(seconds=777),
+                outdated_starting_at=dict(day_of_month=40),
+                repo_path=self.repo_path,
+                settings=dict(mage=dict(partitions=3)),
+            )
+        )
+
+        arr = GlobalDataProduct.load_all(self.repo_path)
         gdp = arr[0]
 
         self.assertEqual(gdp.object_type, 'test1')

mage_ai/tests/data_preparation/models/test_pipeline.py

@@ -725,14 +725,14 @@ class PipelineTest(AsyncDBTestCase):
                     widgets=[],
                 ),
             )
-        pipeline_load = Pipeline.get('test_pipeline_9')
+        pipeline_load = Pipeline.get('test_pipeline_9', repo_path=self.repo_path)
         self.assertEqual(pipeline_load.to_dict()['data_integration'], expected_catalog_config)
 
     def test_save_and_get_integration_pipeline_async(self):
         pipeline = self.__create_pipeline_with_integration('test_pipeline_10')
         asyncio.run(pipeline.save_async())
 
-        pipeline_load = asyncio.run(Pipeline.get_async('test_pipeline_10'))
+        pipeline_load = asyncio.run(Pipeline.get_async('test_pipeline_10', self.repo_path))
         self.assertEqual(
             pipeline_load.to_dict()['data_integration'],
             {
@@ -941,14 +941,17 @@ def export_data(df, *args):
             self.faker.unique.name(),
             repo_path=self.repo_path,
         )
-        self.assertEqual(Pipeline.get(pipeline.uuid).uuid, pipeline.uuid)
+        self.assertEqual(Pipeline.get(pipeline.uuid, repo_path=self.repo_path).uuid, pipeline.uuid)
 
     async def test_get_async(self):
         pipeline = Pipeline.create(
             self.faker.unique.name(),
             repo_path=self.repo_path,
         )
-        self.assertEqual((await Pipeline.get_async(pipeline.uuid)).uuid, pipeline.uuid)
+        self.assertEqual(
+            (await Pipeline.get_async(pipeline.uuid, self.repo_path)).uuid,
+            pipeline.uuid,
+        )
 
     def test_get_all_pipelines_all_projects(self):
         with patch('mage_ai.data_preparation.models.pipeline.Pipeline.get_all_pipelines') as mock:
@@ -1010,7 +1013,7 @@ class PipelineProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
 
         for pipeline in pipelines:
             self.assertIsNone(
-                Pipeline.get(pipeline.uuid, check_if_exists=True, repo_path=base_repo_path()),
+                Pipeline.get(pipeline.uuid, repo_path=base_repo_path(), check_if_exists=True),
             )
 
         with patch(
@@ -1066,8 +1069,8 @@ class PipelineProjectPlatformTests(ProjectPlatformMixin, AsyncDBTestCase):
                 self.assertEqual(
                     (await Pipeline.get_async(
                         pipeline.uuid,
+                        base_repo_path(),
                         all_projects=True,
-                        repo_path=base_repo_path(),
                     )).uuid,
                     pipeline.uuid,
                 )

mage_ai/tests/data_preparation/models/test_project.py

@@ -176,7 +176,10 @@ class ProjectTest(ProjectPlatformMixin, AsyncDBTestCase):
             'mage_ai.data_preparation.models.project.project_platform_settings',
         ) as mock_project_platform_settings:
             Project().projects()
-            mock_project_platform_settings.assert_called_with(mage_projects_only=True)
+            mock_project_platform_settings.assert_called_with(
+                context_data=dict(),
+                mage_projects_only=True,
+            )
 
     def test_is_feature_enabled(self):
         project = Project()

mage_ai/tests/data_preparation/models/test_utils.py

@@ -0,0 +1,80 @@
+from unittest.mock import patch
+
+import pandas as pd
+import polars as pl
+from scipy.sparse import csr_matrix
+from sklearn.linear_model import LinearRegression
+
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
+from mage_ai.data_preparation.models.utils import infer_variable_type
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.tests.base_test import TestCase
+
+
+class TestModelUtils(TestCase):
+    def setUp(self):
+        self.repo_path = self.repo_path
+        self.pipeline_uuid = 'pipeline_uuid'
+        self.block_uuid = 'block_uuid'
+        self.variable_uuid = 'variable_uuid'
+        self.partition = 'partition'
+        self.storage = 'storage'
+        self.clean_block_uuid = 'clean_block_uuid'
+
+    def test_polars_dataframe(self):
+        data = pl.DataFrame({'a': [1, 2, 3]})
+        with patch(
+            'mage_ai.data_preparation.models.project.Project.is_feature_enabled',
+            lambda _, feature_uuid: FeatureUUID.POLARS == feature_uuid,
+        ):
+            variable_type_use, basic_iterable = infer_variable_type(data)
+            self.assertEqual(variable_type_use, VariableType.POLARS_DATAFRAME)
+
+    def test_pandas_dataframe(self):
+        data = pd.DataFrame({'a': [1, 2, 3]})
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.DATAFRAME)
+
+    def test_sparse_matrix(self):
+        data = csr_matrix([[0, 0, 1], [1, 0, 0], [0, 1, 0]])
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.MATRIX_SPARSE)
+
+    def test_pandas_series(self):
+        data = pd.Series([1, 2, 3])
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.SERIES_PANDAS)
+
+    def test_polars_series(self):
+        data = pl.Series([1, 2, 3])
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.SERIES_POLARS)
+
+    def test_sklearn_model(self):
+        data = LinearRegression()
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.MODEL_SKLEARN)
+
+    def test_list_complex(self):
+        data = [TestCase(), TestCase()]
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.LIST_COMPLEX)
+
+    def test_dictionary_complex(self):
+        data = {'a': TestCase(), 'b': TestCase()}
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.DICTIONARY_COMPLEX)
+
+    def test_custom_object(self):
+        class CustomObject:
+            def __init__(self):
+                self.a = 1
+
+        data = CustomObject()
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.CUSTOM_OBJECT)
+
+    def test_iterable(self):
+        data = [1, 2, 3]
+        variable_type_use, basic_iterable = infer_variable_type(data)
+        self.assertEqual(variable_type_use, VariableType.ITERABLE)