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

mage_ai/tests/data_preparation/models/block/dbt/test_block_yaml.py

@@ -1,296 +1,326 @@
-import asyncio
-import os
-import secrets
-from pathlib import Path
-from unittest.mock import MagicMock, patch
+# import asyncio
+# import os
+# import secrets
+# from pathlib import Path
+# from unittest.mock import MagicMock, call, patch
 
-from mage_ai.data_preparation.models.block import Block
-from mage_ai.data_preparation.models.block.dbt.block import DBTBlock
-from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
-from mage_ai.settings.utils import base_repo_path
-from mage_ai.tests.base_test import TestCase
-from mage_ai.tests.data_preparation.models.block.platform.test_mixins import (
-    BlockWithProjectPlatformShared,
-)
-from mage_ai.tests.shared.mixins import ProjectPlatformMixin
+# from dbt.cli.main import dbtRunnerResult
 
+# from mage_ai.data_preparation.models.block import Block
+# from mage_ai.data_preparation.models.block.dbt.block import DBTBlock
+# from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
+# from mage_ai.settings.utils import base_repo_path
+# from mage_ai.tests.base_test import TestCase
+# from mage_ai.tests.data_preparation.models.block.platform.test_mixins import (
+#     BlockWithProjectPlatformShared,
+# )
+# from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
-def build_block(pipeline, content: str) -> DBTBlock:
-    return DBTBlock(
-        name='test_dbt_block_yaml',
-        uuid='test_dbt_block_yaml',
-        block_type=BlockType.DBT,
-        language=BlockLanguage.YAML,
-        pipeline=pipeline,
-        configuration={
-            'dbt_project_name': 'test_project_name',
-            'dbt_profile_target': 'dev',
-            'dbt': {'command': 'build'}
-        },
-        content=content,
-    )
 
+# def build_block(pipeline, content: str) -> DBTBlock:
+#     return DBTBlock(
+#         name='test_dbt_block_yaml',
+#         uuid='test_dbt_block_yaml',
+#         block_type=BlockType.DBT,
+#         language=BlockLanguage.YAML,
+#         pipeline=pipeline,
+#         configuration={
+#             'dbt_project_name': 'test_project_name',
+#             'dbt_profile_target': 'dev',
+#             'dbt': {'command': 'build'}
+#         },
+#         content=content,
+#     )
 
-class DBTBlockYAMLTest(TestCase):
-    @classmethod
-    def setUpClass(self):
-        super().setUpClass()
 
-        self.pipeline = MagicMock()
-        self.pipeline.uuid = 'test'
-        self.pipeline.repo_path = 'test_repo_path'
+# class DBTBlockYAMLTest(TestCase):
+#     @classmethod
+#     def setUpClass(self):
+#         super().setUpClass()
 
-        self.dbt_block = build_block(
-            self.pipeline,
-            '--select model+ --exclude model --vars \'{"foo":"bar"}\'',
-        )
+#         self.pipeline = MagicMock()
+#         self.pipeline.uuid = 'test'
+#         self.pipeline.repo_path = 'test_repo_path'
 
-    @classmethod
-    def tearDownClass(self):
-        super().tearDownClass()
+#         self.dbt_block = build_block(
+#             self.pipeline,
+#             '--select model+ --exclude model --vars \'{"foo":"bar"}\'',
+#         )
 
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Project')
-    def test_metadata_async(self, Project, Profiles):
-        Project.return_value.local_packages = ['test_project_name']
-        Project.return_value.project = {
-            'name': 'test_project_name',
-            'profile': 'test_project_name'
-        }
-        Profiles.return_value.profiles = {
-            'test_project_name': {
-                'target': 'test',
-                'outputs': {
-                    'test': None,
-                    'dev': None,
-                    'prod': None
-                }
-            }
-        }
+#     @classmethod
+#     def tearDownClass(self):
+#         super().tearDownClass()
 
-        metadata = asyncio.run(self.dbt_block.metadata_async())
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Project')
+#     def test_metadata_async(self, Project, Profiles):
+#         Project.return_value.local_packages = ['test_project_name']
+#         Project.return_value.project = {
+#             'name': 'test_project_name',
+#             'profile': 'test_project_name'
+#         }
+#         Profiles.return_value.profiles = {
+#             'test_project_name': {
+#                 'target': 'test',
+#                 'outputs': {
+#                     'test': None,
+#                     'dev': None,
+#                     'prod': None
+#                 }
+#             }
+#         }
 
-        self.assertEqual(
-            metadata,
-            {
-                'dbt': {
-                    'block': {},
-                    'project': None,
-                    'projects': {
-                        'test_project_name': {
-                            'project_name': 'test_project_name',
-                            'target': 'test',
-                            'targets': ['dev', 'prod', 'test']
-                        }
-                    }
-                }
-            }
-        )
+#         metadata = asyncio.run(self.dbt_block.metadata_async())
 
-    def test_tags(self):
-        self.assertEqual(
-            self.dbt_block.tags(),
-            ['build']
-        )
+#         self.assertEqual(
+#             metadata,
+#             {
+#                 'dbt': {
+#                     'block': {},
+#                     'project': None,
+#                     'projects': {
+#                         'test_project_name': {
+#                             'project_name': 'test_project_name',
+#                             'target': 'test',
+#                             'targets': ['dev', 'prod', 'test']
+#                         }
+#                     }
+#                 }
+#             }
+#         )
 
-    def test_project_path(self):
-        self.assertEqual(
-            self.dbt_block.project_path,
-            str(Path('test_repo_path/dbt/test_project_name'))
-        )
+#     def test_tags(self):
+#         self.assertEqual(
+#             self.dbt_block.tags(),
+#             ['build']
+#         )
 
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.DBTCli')
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
-    def test_execute_block(self, Profiles, DBTCli: MagicMock):
-        DBTCli.return_value.invoke.return_value = (None, True)
-        Profiles.return_value.__enter__.return_value.profiles_dir = 'test_profiles_dir'
+#     def test_project_path(self):
+#         self.assertEqual(
+#             self.dbt_block.project_path,
+#             str(Path('test_repo_path/dbt/test_project_name'))
+#         )
 
-        self.dbt_block._execute_block(
-            {},
-            runtime_arguments={
-                '__mage_variables': {
-                    'blocks': {
-                        'test_dbt_block_yaml': {
-                            'configuration': {
-                                'flags': ['--full-refresh']
-                            }
-                        }
-                    }
-                }
-            },
-            global_vars={}
-        )
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.DBTCli.invoke')
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
+#     def test_execute_block(self, Profiles, mock_invoke: MagicMock):
+#         mock_invoke.return_value = dbtRunnerResult(success=True)
+#         Profiles.return_value.__enter__.return_value.profiles_dir = 'test_profiles_dir'
 
-        DBTCli.assert_called_once_with([
-            'build',
-            '--select', 'model+',
-            '--exclude', 'model',
-            '--vars', '{"foo": "bar"}',
-            '--project-dir', str(Path('test_repo_path/dbt/test_project_name')),
-            '--full-refresh',
-            '--target', 'dev',
-            '--profiles-dir', 'test_profiles_dir'
-        ], None)
+#         self.dbt_block._execute_block(
+#             {},
+#             runtime_arguments={
+#                 '__mage_variables': {
+#                     'blocks': {
+#                         'test_dbt_block_yaml': {
+#                             'configuration': {
+#                                 'flags': ['--full-refresh']
+#                             }
+#                         }
+#                     }
+#                 }
+#             },
+#             global_vars={}
+#         )
 
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.DBTCli')
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
-    def test_execute_block_with_interpolation(self, Profiles, DBTCli: MagicMock):
-        DBTCli.return_value.invoke.return_value = (None, True)
-        Profiles.return_value.__enter__.return_value.profiles_dir = 'test_profiles_dir'
+#         self.assertEqual(mock_invoke.mock_calls[0], call([
+#             'deps',
+#             '--select', 'model+',
+#             '--exclude', 'model',
+#             '--vars', '{"foo": "bar"}',
+#             '--project-dir', str(Path('test_repo_path/dbt/test_project_name')),
+#             '--full-refresh',
+#             '--target', 'dev',
+#             '--profiles-dir', 'test_profiles_dir'
+#         ]))
+#         self.assertEqual(mock_invoke.mock_calls[1], call([
+#             'build',
+#             '--select', 'model+',
+#             '--exclude', 'model',
+#             '--vars', '{"foo": "bar"}',
+#             '--project-dir', str(Path('test_repo_path/dbt/test_project_name')),
+#             '--full-refresh',
+#             '--target', 'dev',
+#             '--profiles-dir', 'test_profiles_dir'
+#         ]))
 
-        key = secrets.token_urlsafe()
-        value = secrets.token_urlsafe()
-        os.environ[key] = value
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.DBTCli.invoke')
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
+#     def test_execute_block_with_interpolation(self, Profiles, mock_invoke: MagicMock):
+#         mock_invoke.return_value = dbtRunnerResult(success=True)
+#         Profiles.return_value.__enter__.return_value.profiles_dir = 'test_profiles_dir'
 
-        dbt_block = build_block(
-            self.pipeline,
-            ' '.join([
-                '--select',
-                'models/example/"{} {} {}".sql'.format(
-                    '{' + '{',
-                    "variables('model1')",
-                    '}' + '}',
-                ),
-                'models/example/"{} {} {}".sql'.format(
-                    '{' + '{',
-                    "block_output(parse=lambda arr, _variables: arr[0][0]['model2'])",
-                    '}' + '}',
-                ),
-                '\n',
-                '--vars',
-                (
-                    '\'{}"test1": "{} {} {}", '
-                    '"test2": "{} {} {}", '
-                    '"test3": "{} {} {}", '
-                    '"test4": "{} {} {}"{}\'',
-                )[0].format(
-                    '{',
-                    '{' + '{',
-                    f'env_var("{key}")',
-                    '}' + '}',
-                    '{' + '{',
-                    "variables('model1')",
-                    '}' + '}',
-                    '{' + '{',
-                    'block_output(parse=lambda arr, _variables: arr[1])',
-                    '}' + '}',
-                    '{' + '{',
-                    'block_output("pastel_portal", parse=lambda arr, _variables: arr[2])',
-                    '}' + '}',
-                    '}',
-                ),
-            ])
-        )
+#         key = secrets.token_urlsafe()
+#         value = secrets.token_urlsafe()
+#         os.environ[key] = value
 
-        outputs_from_input_vars = dict(
-            df_1=[dict(model2='my_second_dbt_model')],
-            df_2=[1, 2, 3],
-            fanciful_moon=[dict(model2='my_second_dbt_model')],
-            pastel_portal=[1, 2, 3],
-        )
+#         dbt_block = build_block(
+#             self.pipeline,
+#             ' '.join([
+#                 '--select',
+#                 'models/example/"{} {} {}".sql'.format(
+#                     '{' + '{',
+#                     "variables('model1')",
+#                     '}' + '}',
+#                 ),
+#                 'models/example/"{} {} {}".sql'.format(
+#                     '{' + '{',
+#                     "block_output(parse=lambda arr, _variables: arr[0][0]['model2'])",
+#                     '}' + '}',
+#                 ),
+#                 '\n',
+#                 '--vars',
+#                 (
+#                     '\'{}"test1": "{} {} {}", '
+#                     '"test2": "{} {} {}", '
+#                     '"test3": "{} {} {}", '
+#                     '"test4": "{} {} {}"{}\'',
+#                 )[0].format(
+#                     '{',
+#                     '{' + '{',
+#                     f'env_var("{key}")',
+#                     '}' + '}',
+#                     '{' + '{',
+#                     "variables('model1')",
+#                     '}' + '}',
+#                     '{' + '{',
+#                     'block_output(parse=lambda arr, _variables: arr[1])',
+#                     '}' + '}',
+#                     '{' + '{',
+#                     'block_output("pastel_portal", parse=lambda arr, _variables: arr[2])',
+#                     '}' + '}',
+#                     '}',
+#                 ),
+#             ])
+#         )
 
-        dbt_block.upstream_blocks = [
-            Block(
-                'fanciful_moon',
-                'fanciful_moon',
-                BlockType.DATA_LOADER,
-            ),
-            Block(
-                'pastel_portal',
-                'pastel_portal',
-                BlockType.DATA_LOADER,
-            ),
-        ]
+#         outputs_from_input_vars = dict(
+#             df_1=[dict(model2='my_second_dbt_model')],
+#             df_2=[1, 2, 3],
+#             fanciful_moon=[dict(model2='my_second_dbt_model')],
+#             pastel_portal=[1, 2, 3],
+#         )
 
-        dbt_block._execute_block(
-            outputs_from_input_vars,
-            runtime_arguments={
-                '__mage_variables': {
-                    'blocks': {
-                        'test_dbt_block_yaml': {
-                            'configuration': {
-                                'flags': ['--full-refresh']
-                            }
-                        }
-                    }
-                }
-            },
-            global_vars=dict(
-                foo='bar',
-                model1='my_first_dbt_model',
-            ),
-        )
+#         dbt_block.upstream_blocks = [
+#             Block(
+#                 'fanciful_moon',
+#                 'fanciful_moon',
+#                 BlockType.DATA_LOADER,
+#             ),
+#             Block(
+#                 'pastel_portal',
+#                 'pastel_portal',
+#                 BlockType.DATA_LOADER,
+#             ),
+#         ]
 
-        DBTCli.assert_called_once_with([
-            'build',
-            '--select',
-            'models/example/my_first_dbt_model.sql',
-            'models/example/my_second_dbt_model.sql',
-            '--vars',
-            '{"foo": "bar", "model1": "my_first_dbt_model", '
-            f'"test1": "{value}", "test2": "my_first_dbt_model", '
-            '"test3": "[1, 2, 3]", "test4": "3"}',
-            '--project-dir',
-            str(Path('test_repo_path/dbt/test_project_name')),
-            '--full-refresh',
-            '--target',
-            'dev',
-            '--profiles-dir',
-            'test_profiles_dir',
-        ], None)
+#         dbt_block._execute_block(
+#             outputs_from_input_vars,
+#             runtime_arguments={
+#                 '__mage_variables': {
+#                     'blocks': {
+#                         'test_dbt_block_yaml': {
+#                             'configuration': {
+#                                 'flags': ['--full-refresh']
+#                             }
+#                         }
+#                     }
+#                 }
+#             },
+#             global_vars=dict(
+#                 foo='bar',
+#                 model1='my_first_dbt_model',
+#             ),
+#         )
 
+#         self.assertEqual(mock_invoke.mock_calls[0], call([
+#             'deps',
+#             '--select',
+#             'models/example/my_first_dbt_model.sql',
+#             'models/example/my_second_dbt_model.sql',
+#             '--vars',
+#             '{"foo": "bar", "model1": "my_first_dbt_model", '
+#             f'"test1": "{value}", "test2": "my_first_dbt_model", '
+#             '"test3": "[1, 2, 3]", "test4": "3"}',
+#             '--project-dir',
+#             str(Path('test_repo_path/dbt/test_project_name')),
+#             '--full-refresh',
+#             '--target',
+#             'dev',
+#             '--profiles-dir',
+#             'test_profiles_dir',
+#         ]))
 
-@patch(
-    'mage_ai.data_preparation.models.block.platform.mixins.project_platform_activated',
-    lambda: True,
-)
-@patch(
-    'mage_ai.data_preparation.models.block.platform.utils.project_platform_activated',
-    lambda: True,
-)
-@patch('mage_ai.settings.platform.project_platform_activated', lambda: True)
-class DBTBlockYAMLProjectPlatformTest(ProjectPlatformMixin, BlockWithProjectPlatformShared):
-    def test_project_path(self):
-        block = build_block(MagicMock(), '')
-        block.configuration['dbt_project_name'] = 'mage_data/dbt/demo'
-        self.assertEqual(block.project_path, os.path.join(base_repo_path(), 'mage_data/dbt/demo'))
+#         self.assertEqual(mock_invoke.mock_calls[1], call([
+#             'build',
+#             '--select',
+#             'models/example/my_first_dbt_model.sql',
+#             'models/example/my_second_dbt_model.sql',
+#             '--vars',
+#             '{"foo": "bar", "model1": "my_first_dbt_model", '
+#             f'"test1": "{value}", "test2": "my_first_dbt_model", '
+#             '"test3": "[1, 2, 3]", "test4": "3"}',
+#             '--project-dir',
+#             str(Path('test_repo_path/dbt/test_project_name')),
+#             '--full-refresh',
+#             '--target',
+#             'dev',
+#             '--profiles-dir',
+#             'test_profiles_dir',
+#         ]))
 
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
-    @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Project')
-    def test_metadata_async(self, Project, Profiles):
-        Project.return_value.local_packages = ['test_project_name']
-        Project.return_value.project = {
-            'name': 'test_project_name',
-            'profile': 'test_project_name'
-        }
-        Profiles.return_value.profiles = {
-            'test_project_name': {
-                'target': 'test',
-                'outputs': {
-                    'test': None,
-                    'dev': None,
-                    'prod': None
-                }
-            }
-        }
 
-        block = build_block(MagicMock(), '')
-        # block.configuration['dbt_project_name'] = 'demo'
-        metadata = asyncio.run(block.metadata_async())
+# @patch(
+#     'mage_ai.data_preparation.models.block.platform.mixins.project_platform_activated',
+#     lambda: True,
+# )
+# @patch(
+#     'mage_ai.data_preparation.models.block.platform.utils.project_platform_activated',
+#     lambda: True,
+# )
+# @patch('mage_ai.settings.platform.project_platform_activated', lambda: True)
+# class DBTBlockYAMLProjectPlatformTest(ProjectPlatformMixin, BlockWithProjectPlatformShared):
+#     def test_project_path(self):
+#         block = build_block(MagicMock(), '')
+#         block.configuration['dbt_project_name'] = 'mage_data/dbt/demo'
+#         self.assertEqual(block.project_path, os.path.join(base_repo_path(), 'mage_data/dbt/demo'))
 
-        self.assertEqual(
-            metadata,
-            {
-                'dbt': {
-                    'block': {},
-                    'project': None,
-                    'projects': {
-                        'test_project_name': {
-                            'project_name': 'test_project_name',
-                            'target': 'test',
-                            'targets': ['dev', 'prod', 'test']
-                        }
-                    }
-                }
-            }
-        )
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Profiles')
+#     @patch('mage_ai.data_preparation.models.block.dbt.block_yaml.Project')
+#     def test_metadata_async(self, Project, Profiles):
+#         Project.return_value.local_packages = ['test_project_name']
+#         Project.return_value.project = {
+#             'name': 'test_project_name',
+#             'profile': 'test_project_name'
+#         }
+#         Profiles.return_value.profiles = {
+#             'test_project_name': {
+#                 'target': 'test',
+#                 'outputs': {
+#                     'test': None,
+#                     'dev': None,
+#                     'prod': None
+#                 }
+#             }
+#         }
+
+#         block = build_block(MagicMock(), '')
+#         # block.configuration['dbt_project_name'] = 'demo'
+#         metadata = asyncio.run(block.metadata_async())
+
+#         self.assertEqual(
+#             metadata,
+#             {
+#                 'dbt': {
+#                     'block': {},
+#                     'project': None,
+#                     'projects': {
+#                         'test_project_name': {
+#                             'project_name': 'test_project_name',
+#                             'target': 'test',
+#                             'targets': ['dev', 'prod', 'test']
+#                         }
+#                     }
+#                 }
+#             }
+#         )

mage_ai/tests/data_preparation/models/block/dbt/test_dbt_adapter.py

@@ -7,10 +7,10 @@ from unittest.mock import patch
 from dbt.include.starter_project import PACKAGE_PATH as starter_project_directory
 
 from mage_ai.data_preparation.models.block.dbt.dbt_adapter import DBTAdapter
-from mage_ai.tests.base_test import TestCase
+from mage_ai.tests.base_test import AsyncDBTestCase
 
 
-class DBTAdapterTest(TestCase):
+class DBTAdapterTest(AsyncDBTestCase):
     """
     Tests the Project Class, which is an interface with dbt dbt_project.yml files
     """
@@ -114,10 +114,17 @@ class DBTAdapterTest(TestCase):
                     agate_table=table
                 )
 
+                csv_file_path = os.path.join(self.repo_path, 'test_dbt_adapter.csv')
+                df.to_csv(csv_file_path, index=False)
+
                 dbt_adapter.execute_macro(
                     'load_csv_rows',
                     context_overide=relation_context,
-                    model={'config': {}},
+                    model={
+                        'config': {},
+                        'original_file_path': 'test_dbt_adapter.csv',
+                        'root_path': self.repo_path,
+                    },
                     agate_table=table
                 )
 
@@ -125,7 +132,10 @@ class DBTAdapterTest(TestCase):
 
                 self.assertEqual(
                     df.to_dict(),
-                    {0: {'id': Decimal('1'), 'text': Decimal('2')}, 1: {'id': 'foo', 'text': 'bar'}}
+                    {
+                        0: {'id': Decimal('1'), 'text': Decimal('2')},
+                        1: {'id': 'foo', 'text': 'bar'},
+                    },
                 )
 
     def test_credentials(self):