mage-ai 0.9.73__py3-none-any.whl → 0.9.75__py3-none-any.whl

mage_ai/tests/api/endpoints/test_dbt_blocks.py

@@ -0,0 +1,111 @@
+import os
+
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.tests.api.endpoints.mixins import (
+    BaseAPIEndpointTest,
+    build_create_endpoint_tests,
+)
+from mage_ai.tests.shared.mixins import DBTMixin
+
+
+def get_pipeline(self):
+    return Pipeline.get(self.pipeline.uuid, repo_path=get_repo_path())
+
+
+class DBTBlockAPIEndpointTest(BaseAPIEndpointTest, DBTMixin):
+    pass
+
+
+def __assert_after_create(self):
+    pipeline = get_pipeline(self)
+    blocks = pipeline.blocks_by_uuid.values()
+    blocks_dbt = list(filter(
+        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.SQL,
+        blocks,
+    ))
+
+    self.assertEqual(len(blocks), 5)
+    self.assertEqual(len(blocks_dbt), 1)
+
+    block = blocks_dbt[0]
+    self.assertEqual({
+      'dbt_project_name': 'dbt/test_project',
+      'file_path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
+      'file_source': {
+        'path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
+        'project_path': 'dbt/test_project',
+      },
+    }, block.configuration)
+
+    return True
+
+
+def __assert_after_create_yaml(self):
+    pipeline = get_pipeline(self)
+    blocks = pipeline.blocks_by_uuid.values()
+    blocks_dbt = list(filter(
+        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.YAML,
+        blocks,
+    ))
+
+    self.assertEqual(len(blocks), 5)
+    self.assertEqual(len(blocks_dbt), 1)
+
+    block = blocks_dbt[0]
+    self.assertEqual({
+      'dbt_project_name': 'dbt/test_project',
+      'file_source': {
+        'path': 'dbts/dbt_yaml.yaml',
+        'project_path': 'dbt/test_project',
+      },
+    }, block.configuration)
+
+    return True
+
+
+build_create_endpoint_tests(
+    DBTBlockAPIEndpointTest,
+    assert_after_create_count=__assert_after_create,
+    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
+    build_payload=lambda self: dict(
+        name=self.faker.unique.name(),
+        type=BlockType.DBT,
+        language=BlockLanguage.SQL,
+        configuration=dict(
+            file_path=os.path.join(
+                os.path.basename(self.dbt_directory),
+                'test_project',
+                'models',
+                'example',
+                'my_first_dbt_model.sql',
+            ),
+        ),
+    ),
+    get_resource_parent_id=lambda self: self.pipeline.uuid,
+    resource='block',
+    resource_parent='pipelines',
+)
+
+
+build_create_endpoint_tests(
+    DBTBlockAPIEndpointTest,
+    test_uuid='dbt_yaml',
+    assert_after_create_count=__assert_after_create_yaml,
+    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
+    build_payload=lambda self: dict(
+        name='dbt_yaml',
+        type=BlockType.DBT,
+        language=BlockLanguage.YAML,
+        configuration=dict(
+            dbt_project_name=os.path.join(
+                os.path.basename(self.dbt_directory),
+                'test_project',
+            ),
+        ),
+    ),
+    get_resource_parent_id=lambda self: self.pipeline.uuid,
+    resource='block',
+    resource_parent='pipelines',
+)

mage_ai/tests/api/endpoints/test_file_contents.py

@@ -1,12 +1,10 @@
 import urllib.parse
-from unittest.mock import patch
 
 from mage_ai.tests.api.endpoints.mixins import (
     BaseAPIEndpointTest,
     build_detail_endpoint_tests,
     build_update_endpoint_tests,
 )
-from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
 
 def get_resource_id(self) -> str:
@@ -44,49 +42,3 @@ build_update_endpoint_tests(
     get_model_before_update=get_model_before_update,
     assert_after_update=lambda self, result, model: model != result['content'],
 )
-
-
-class FileContentWithProjectPlatformAPIEndpointTest(BaseAPIEndpointTest, ProjectPlatformMixin):
-    def setUp(self):
-        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
-            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
-                super().setUp()
-                self.setup_final()
-
-    def tearDown(self):
-        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
-            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
-                self.teardown_final()
-                super().tearDown()
-
-
-build_detail_endpoint_tests(
-    FileContentWithProjectPlatformAPIEndpointTest,
-    resource='file_content',
-    get_resource_id=get_resource_id,
-    result_keys_to_compare=[
-        'content',
-        'name',
-        'path',
-    ],
-    patch_function_settings=[
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
-        ('mage_ai.api.resources.FileContentResource.ensure_file_is_in_project', lambda _x: True),
-    ],
-)
-
-
-build_update_endpoint_tests(
-    FileContentWithProjectPlatformAPIEndpointTest,
-    resource='file_content',
-    get_resource_id=get_resource_id,
-    build_payload=lambda self: dict(content=self.faker.unique.name()),
-    get_model_before_update=get_model_before_update,
-    assert_after_update=lambda self, result, model, mocks: model != result['content'],
-    patch_function_settings=[
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
-        ('mage_ai.api.resources.FileContentResource.ensure_file_is_in_project', lambda _x: True),
-    ],
-)

mage_ai/tests/api/endpoints/test_file_contents_with_project_platform.py

@@ -0,0 +1,66 @@
+import urllib.parse
+from unittest.mock import patch
+
+from mage_ai.tests.api.endpoints.mixins import (
+    BaseAPIEndpointTest,
+    build_detail_endpoint_tests,
+    build_update_endpoint_tests,
+)
+from mage_ai.tests.shared.mixins import ProjectPlatformMixin
+
+
+def get_resource_id(self) -> str:
+    block = list(self.pipeline.blocks_by_uuid.values())[0]
+    pk = block.file_path.replace(f'{self.pipeline.repo_path}/', '')
+    return urllib.parse.quote_plus(pk)
+
+
+def get_model_before_update(self):
+    block = list(self.pipeline.blocks_by_uuid.values())[0]
+    return block.file.content()
+
+
+class FileContentWithProjectPlatformAPIEndpointTest(BaseAPIEndpointTest, ProjectPlatformMixin):
+    def setUp(self):
+        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
+            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
+                super().setUp()
+                self.setup_final()
+
+    def tearDown(self):
+        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
+            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
+                self.teardown_final()
+                super().tearDown()
+
+
+build_detail_endpoint_tests(
+    FileContentWithProjectPlatformAPIEndpointTest,
+    resource='file_content',
+    get_resource_id=get_resource_id,
+    result_keys_to_compare=[
+        'content',
+        'name',
+        'path',
+    ],
+    patch_function_settings=[
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
+        ('mage_ai.api.resources.FileContentResource.ensure_file_is_in_project', lambda _x: True),
+    ],
+)
+
+
+build_update_endpoint_tests(
+    FileContentWithProjectPlatformAPIEndpointTest,
+    resource='file_content',
+    get_resource_id=get_resource_id,
+    build_payload=lambda self: dict(content=self.faker.unique.name()),
+    get_model_before_update=get_model_before_update,
+    assert_after_update=lambda self, result, model, mocks: model != result['content'],
+    patch_function_settings=[
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
+        ('mage_ai.api.resources.FileContentResource.ensure_file_is_in_project', lambda _x: True),
+    ],
+)

mage_ai/tests/api/endpoints/test_pipelines.py

@@ -1,5 +1,3 @@
-from unittest.mock import patch
-
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.tests.api.endpoints.mixins import (
     BaseAPIEndpointTest,
@@ -9,7 +7,6 @@ from mage_ai.tests.api.endpoints.mixins import (
     build_list_endpoint_tests,
     build_update_endpoint_tests,
 )
-from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
 
 class PipelineAPIEndpointTest(BaseAPIEndpointTest):
@@ -111,134 +108,3 @@ build_delete_endpoint_tests(
         Pipeline.get_all_pipelines(self.repo_path),
     ) == 1,
 )
-
-
-class PipelineWithProjectPlatformAPIEndpointTest(BaseAPIEndpointTest, ProjectPlatformMixin):
-    def setUp(self):
-        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
-            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
-                super().setUp()
-                self.setup_final()
-
-    def tearDown(self):
-        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
-            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
-                self.teardown_final()
-                super().tearDown()
-
-
-build_list_endpoint_tests(
-    PipelineWithProjectPlatformAPIEndpointTest,
-    resource='pipelines',
-    list_count=3,
-    result_keys_to_compare=[
-        'blocks',
-        'callbacks',
-        'concurrency_config',
-        'conditionals',
-        'created_at',
-        'data_integration',
-        'description',
-        'executor_config',
-        'executor_count',
-        'executor_type',
-        'name',
-        'notification_config',
-        'retry_config',
-        'run_pipeline_in_one_process',
-        'spark_config',
-        'tags',
-        'type',
-        'updated_at',
-        'uuid',
-        'widgets',
-    ],
-    patch_function_settings=[
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
-    ],
-)
-
-
-build_create_endpoint_tests(
-    PipelineWithProjectPlatformAPIEndpointTest,
-    resource='pipelines',
-    assert_after_create_count=lambda self, mocks: len(
-        Pipeline.get_all_pipelines(self.repo_path),
-    ) == 4,
-    assert_before_create_count=lambda self, mocks: len(
-        Pipeline.get_all_pipelines(self.repo_path),
-    ) == 3,
-    build_payload=lambda self: dict(
-        name=self.faker.unique.name(),
-    ),
-    patch_function_settings=[
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
-    ],
-)
-
-
-build_detail_endpoint_tests(
-    PipelineWithProjectPlatformAPIEndpointTest,
-    resource='pipelines',
-    get_resource_id=lambda self: self.pipeline.uuid,
-    result_keys_to_compare=[
-        'blocks',
-        'callbacks',
-        'concurrency_config',
-        'conditionals',
-        'created_at',
-        'data_integration',
-        'description',
-        'executor_config',
-        'executor_count',
-        'executor_type',
-        'extensions',
-        'name',
-        'notification_config',
-        'retry_config',
-        'run_pipeline_in_one_process',
-        'spark_config',
-        'tags',
-        'type',
-        'updated_at',
-        'uuid',
-        'widgets',
-    ],
-    patch_function_settings=[
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
-    ],
-)
-
-
-build_update_endpoint_tests(
-    PipelineWithProjectPlatformAPIEndpointTest,
-    resource='pipelines',
-    get_resource_id=lambda self: self.pipeline.uuid,
-    build_payload=lambda self: dict(name=self.faker.unique.name()),
-    get_model_before_update=lambda self: self.pipeline,
-    assert_after_update=lambda self, result, model, mocks: model.name != result['name'],
-    patch_function_settings=[
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
-    ],
-)
-
-
-build_delete_endpoint_tests(
-    PipelineWithProjectPlatformAPIEndpointTest,
-    resource='pipelines',
-    get_resource_id=lambda self: self.pipeline.uuid,
-    assert_after_delete_count=lambda self, mocks: len(
-        Pipeline.get_all_pipelines(self.repo_path),
-    ) == 2,
-    assert_before_delete_count=lambda self, mocks: len(
-        Pipeline.get_all_pipelines(self.repo_path),
-    ) == 3,
-    patch_function_settings=[
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
-    ],
-)

mage_ai/tests/api/endpoints/test_pipelines_with_project_platform.py

@@ -0,0 +1,143 @@
+from unittest.mock import patch
+
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.tests.api.endpoints.mixins import (
+    BaseAPIEndpointTest,
+    build_create_endpoint_tests,
+    build_delete_endpoint_tests,
+    build_detail_endpoint_tests,
+    build_list_endpoint_tests,
+    build_update_endpoint_tests,
+)
+from mage_ai.tests.shared.mixins import ProjectPlatformMixin
+
+
+class PipelineWithProjectPlatformAPIEndpointTest(BaseAPIEndpointTest, ProjectPlatformMixin):
+    def setUp(self):
+        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
+            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
+                super().setUp()
+                self.setup_final()
+
+    def tearDown(self):
+        with patch('mage_ai.settings.platform.project_platform_activated', lambda: True):
+            with patch('mage_ai.settings.repo.project_platform_activated', lambda: True):
+                self.teardown_final()
+                super().tearDown()
+
+
+build_list_endpoint_tests(
+    PipelineWithProjectPlatformAPIEndpointTest,
+    resource='pipelines',
+    list_count=3,
+    result_keys_to_compare=[
+        'blocks',
+        'callbacks',
+        'concurrency_config',
+        'conditionals',
+        'created_at',
+        'data_integration',
+        'description',
+        'executor_config',
+        'executor_count',
+        'executor_type',
+        'name',
+        'notification_config',
+        'retry_config',
+        'run_pipeline_in_one_process',
+        'spark_config',
+        'tags',
+        'type',
+        'updated_at',
+        'uuid',
+        'widgets',
+    ],
+    patch_function_settings=[
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
+    ],
+)
+
+
+build_create_endpoint_tests(
+    PipelineWithProjectPlatformAPIEndpointTest,
+    resource='pipelines',
+    assert_after_create_count=lambda self, mocks: len(
+        Pipeline.get_all_pipelines(self.repo_path),
+    ) == 4,
+    assert_before_create_count=lambda self, mocks: len(
+        Pipeline.get_all_pipelines(self.repo_path),
+    ) == 3,
+    build_payload=lambda self: dict(
+        name=self.faker.unique.name(),
+    ),
+    patch_function_settings=[
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
+    ],
+)
+
+
+build_detail_endpoint_tests(
+    PipelineWithProjectPlatformAPIEndpointTest,
+    resource='pipelines',
+    get_resource_id=lambda self: self.pipeline.uuid,
+    result_keys_to_compare=[
+        'blocks',
+        'callbacks',
+        'concurrency_config',
+        'conditionals',
+        'created_at',
+        'data_integration',
+        'description',
+        'executor_config',
+        'executor_count',
+        'executor_type',
+        'extensions',
+        'name',
+        'notification_config',
+        'retry_config',
+        'run_pipeline_in_one_process',
+        'spark_config',
+        'tags',
+        'type',
+        'updated_at',
+        'uuid',
+        'widgets',
+    ],
+    patch_function_settings=[
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
+    ],
+)
+
+
+build_update_endpoint_tests(
+    PipelineWithProjectPlatformAPIEndpointTest,
+    resource='pipelines',
+    get_resource_id=lambda self: self.pipeline.uuid,
+    build_payload=lambda self: dict(name=self.faker.unique.name()),
+    get_model_before_update=lambda self: self.pipeline,
+    assert_after_update=lambda self, result, model, mocks: model.name != result['name'],
+    patch_function_settings=[
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
+    ],
+)
+
+
+build_delete_endpoint_tests(
+    PipelineWithProjectPlatformAPIEndpointTest,
+    resource='pipelines',
+    get_resource_id=lambda self: self.pipeline.uuid,
+    assert_after_delete_count=lambda self, mocks: len(
+        Pipeline.get_all_pipelines(self.repo_path),
+    ) == 2,
+    assert_before_delete_count=lambda self, mocks: len(
+        Pipeline.get_all_pipelines(self.repo_path),
+    ) == 3,
+    patch_function_settings=[
+        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
+        ('mage_ai.settings.platform.project_platform_activated', lambda: True),
+    ],
+)

mage_ai/tests/api/policies/permissions/test_base_policy_with_permissions.py

@@ -1,5 +1,4 @@
 import secrets
-from enum import Enum
 from typing import Any, Callable, Dict, List, Union
 from unittest.mock import patch
 
@@ -19,13 +18,14 @@ from mage_ai.authentication.permissions.constants import (
 )
 from mage_ai.orchestration.db.models.oauth import Permission, Role, User, UserRole
 from mage_ai.shared.array import find
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import index_by, merge_dict
 from mage_ai.tests.api.mixins import BootstrapMixin
 from mage_ai.tests.api.operations.test_base import BaseApiTestCase
 from mage_ai.tests.api.policies.permissions.mixins import PermissionsMixin
 
 
-class TestSuite(str, Enum):
+class TestSuite(StrEnum):
     AUTHORIZED = 'AUTHORIZED'
     DISABLED = 'DISABLED'
     INVERSE = 'INVERSE'