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

mage_ai/tests/data_preparation/executors/test_block_executor.py

@@ -15,6 +15,7 @@ from mage_ai.data_preparation.models.block.hook.block import HookBlock
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun
+from mage_ai.shared.hash import ignore_keys, merge_dict
 from mage_ai.tests.api.operations.test_base import BaseApiTestCase
 from mage_ai.tests.factory import create_pipeline_with_blocks
 from mage_ai.tests.shared.mixins import build_hooks
@@ -33,6 +34,7 @@ class BlockExecutorTest(BaseApiTestCase):
         self.block.configuration = {}
         self.block.upstream_blocks = []
         self.block.type = BlockType.DBT
+        self.block.replicated_block = None
 
         self.pipeline.get_block.return_value = self.block
         self.pipeline.repo_config.retry_config = {'retries': 3, 'delay': 1}
@@ -75,7 +77,9 @@ class BlockExecutorTest(BaseApiTestCase):
     def test_execute(self):
         analyze_outputs = False
         callback_url = 'http://example.com/callback'
-        global_vars = {'var1': 'value1'}
+        global_vars = {
+            'var1': 'value1',
+        }
         update_status = True
 
         def on_complete(block_uuid, metrics: Dict = None):
@@ -133,7 +137,7 @@ class BlockExecutorTest(BaseApiTestCase):
             block_run_outputs_cache=None,
             cache_block_output_in_memory=False,
             callback_url=callback_url,
-            global_vars=global_vars,
+            global_vars=merge_dict(global_vars, dict(retry=dict(attempts=1))),
             input_from_output=input_from_output,
             logging_tags={
                 'block_type': BlockType.DBT,
@@ -167,9 +171,10 @@ class BlockExecutorTest(BaseApiTestCase):
         # )
         self.block_executor._execute_callback.assert_called_with(
             'on_success',
+            callback_kwargs=dict(retry=dict(attempts=1)),
             dynamic_block_index=None,
             dynamic_upstream_block_uuids=None,
-            global_vars=global_vars,
+            global_vars=ignore_keys(global_vars, ['retry']),
             logging_tags={
                 'block_type': BlockType.DBT,
                 'block_uuid': self.block_uuid,
@@ -243,7 +248,7 @@ class BlockExecutorTest(BaseApiTestCase):
 
         self.block_executor._execute_callback(
             callback='on_success',
-            global_vars={},
+            global_vars=dict(retry=dict(attempts=1)),
             logging_tags={},
             pipeline_run=None,
             dynamic_block_index=None,
@@ -255,7 +260,7 @@ class BlockExecutorTest(BaseApiTestCase):
             dynamic_block_index=None,
             dynamic_upstream_block_uuids=None,
             execution_partition=self.execution_partition,
-            global_vars={},
+            global_vars=dict(retry=dict(attempts=1)),
             logger=self.logger,
             logging_tags={},
             parent_block=self.block,
@@ -374,15 +379,21 @@ class BlockExecutorTest(BaseApiTestCase):
                 executor.execute(block_run_id=block_run.id, on_complete=__on_complete)
 
     def test_block_run_for_dynamic_child_block(self):
-        block1 = self.blocks[0]
+        pipeline, blocks = create_pipeline_with_blocks(
+            'dynamic_child_block',
+            self.repo_path,
+            return_blocks=True,
+        )
+
+        block1 = blocks[0]
         block1.configuration = dict(dynamic=True)
-        self.pipeline1.add_block(block1)
-        block2 = self.blocks[1]
+        pipeline.add_block(block1)
+        block2 = blocks[1]
 
         pipeline_run = PipelineRun.create(
             execution_date=datetime.utcnow(),
             pipeline_schedule_id=0,
-            pipeline_uuid=self.pipeline1.uuid,
+            pipeline_uuid=pipeline.uuid,
         )
         block_run = BlockRun.create(
             block_uuid=block2.uuid,
@@ -390,20 +401,18 @@ class BlockExecutorTest(BaseApiTestCase):
         )
 
         executor = BlockExecutor(
-            self.pipeline1,
+            pipeline,
             block_run.block_uuid,
+            block_run_id=block_run.id,
             execution_partition=pipeline_run.execution_partition,
         )
 
         self.assertTrue(isinstance(executor.block, DynamicChildBlockFactory))
-        self.assertTrue(executor.block.pipeline_run is None)
 
         with patch.object(executor.block, 'execute_sync') as mock_execute_sync:
             executor.execute(block_run_id=block_run.id, pipeline_run_id=pipeline_run.id)
             mock_execute_sync.assert_called_once()
 
-            self.assertEqual(executor.block.pipeline_run.id, pipeline_run.id)
-
         block_run_dynamic_child_spawn = BlockRun.create(
             block_uuid=f'{block2.uuid}:0',
             pipeline_run_id=pipeline_run.id,
@@ -412,46 +421,47 @@ class BlockExecutorTest(BaseApiTestCase):
             }))
         )
         executor = BlockExecutor(
-            self.pipeline1,
+            pipeline,
             block_run_dynamic_child_spawn.block_uuid,
             execution_partition=pipeline_run.execution_partition,
         )
 
         self.assertFalse(isinstance(executor.block, DynamicChildBlockFactory))
 
-        with patch.object(self.pipeline1, 'get_block', lambda _x: block2):
-            with patch.object(block2, 'execute_sync') as mock_execute_sync:
-                executor.execute(
+        # with patch.object(pipeline, 'get_block', lambda _x: block2):
+        with patch.object(block2, 'execute_sync') as mock_execute_sync:
+            executor.execute(
+                block_run_id=block_run_dynamic_child_spawn.id,
+                pipeline_run_id=pipeline_run.id,
+            )
+
+            mock_execute_sync.assert_called_once_with(
+                analyze_outputs=False,
+                block_run_outputs_cache=None,
+                dynamic_block_index=None,
+                dynamic_block_indexes=dict({
+                    block1.uuid: 0,
+                }),
+                dynamic_block_uuid=None,
+                dynamic_upstream_block_uuids=None,
+                execution_partition=pipeline_run.execution_partition,
+                global_vars=dict(retry=dict(attempts=1), logger=ANY),
+                input_from_output=None,
+                logger=ANY,
+                logging_tags=dict(
                     block_run_id=block_run_dynamic_child_spawn.id,
+                    block_type=block2.type,
+                    block_uuid=block_run_dynamic_child_spawn.block_uuid,
                     pipeline_run_id=pipeline_run.id,
-                )
-
-                mock_execute_sync.assert_called_once_with(
-                    analyze_outputs=False,
-                    block_run_outputs_cache=None,
-                    dynamic_block_index=None,
-                    dynamic_block_indexes=dict({
-                        block1.uuid: 0,
-                    }),
-                    dynamic_block_uuid=None,
-                    dynamic_upstream_block_uuids=None,
-                    execution_partition=pipeline_run.execution_partition,
-                    global_vars=None,
-                    input_from_output=None,
-                    logger=ANY,
-                    logging_tags=dict(
-                        block_run_id=block_run_dynamic_child_spawn.id,
-                        block_type=block2.type,
-                        block_uuid=block_run_dynamic_child_spawn.block_uuid,
-                        pipeline_run_id=pipeline_run.id,
-                        pipeline_uuid=self.pipeline1.uuid,
-                    ),
-                    run_all_blocks=True,
-                    runtime_arguments=None,
-                    store_variables=True,
-                    update_status=False,
-                    verify_output=True,
-                )
+                    pipeline_uuid=pipeline.uuid,
+                ),
+                metadata=None,
+                run_all_blocks=True,
+                runtime_arguments=None,
+                store_variables=True,
+                update_status=False,
+                verify_output=True,
+            )
 
     def test_block_run_for_dynamic_child_block_reduce_output(self):
         block1 = self.blocks[0]

mage_ai/tests/data_preparation/git/test_api.py

@@ -3,7 +3,7 @@ from datetime import timedelta
 from faker import Faker
 
 from mage_ai.authentication.oauth2 import generate_access_token
-from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_GITHUB
+from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.data_preparation.git.api import (
     get_access_token_for_user,
     get_oauth_client_id,
@@ -22,12 +22,12 @@ class GitApiTest(DBTestCase):
         email = self.faker.email()
         user = create_user(email=email, roles=1)
 
-        client_id = get_oauth_client_id(OAUTH_PROVIDER_GITHUB)
+        client_id = get_oauth_client_id(ProviderName.GITHUB)
 
         oauth_client = Oauth2Application.create(
             client_id=client_id,
             client_type=Oauth2Application.ClientType.PRIVATE,
-            name=OAUTH_PROVIDER_GITHUB,
+            name=ProviderName.GITHUB,
             user_id=user.id if user else None,
         )
 
@@ -37,5 +37,5 @@ class GitApiTest(DBTestCase):
             duration=int(timedelta(days=30).total_seconds()),
         )
 
-        access_token = get_access_token_for_user(user, provider=OAUTH_PROVIDER_GITHUB)
+        access_token = get_access_token_for_user(user, provider=ProviderName.GITHUB)
         self.assertIsNotNone(access_token)

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

@@ -1,21 +1,35 @@
-from pathlib import Path
-from unittest.mock import MagicMock, patch
+import os
+from unittest.mock import MagicMock
 
 from mage_ai.data_preparation.models.block.dbt import DBTBlock
 from mage_ai.data_preparation.models.block.dbt.block_sql import DBTBlockSQL
 from mage_ai.data_preparation.models.block.dbt.block_yaml import DBTBlockYAML
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
-from mage_ai.tests.base_test import TestCase
+from mage_ai.tests.base_test import AsyncDBTestCase
 
 
-class DBTBlockTest(TestCase):
+class DBTBlockTest(AsyncDBTestCase):
     @classmethod
     def setUpClass(self):
         super().setUpClass()
 
         pipeline = MagicMock()
         pipeline.uuid = 'test'
-        pipeline.repo_path = 'test_repo_path'
+        pipeline.repo_path = self.repo_path
+
+        project_path = os.path.join(self.repo_path, 'dbt', 'dir1', 'dir2', 'dir3', 'fire')
+
+        for filename in [
+            'profiles.yml',
+            'dbt_project.yml',
+            os.path.join('models', 'demo', 'model.sql')
+        ]:
+            fp = os.path.join(project_path, filename)
+            os.makedirs(os.path.dirname(fp), exist_ok=True)
+            with open(fp, 'w') as f:
+                f.write('')
+
+        os.makedirs(os.path.join(self.repo_path, 'dbts'), exist_ok=True)
 
         self.block_sql = MagicMock()
         self.block_sql.name = 'test_block_sql'
@@ -30,7 +44,7 @@ class DBTBlockTest(TestCase):
             block_type=BlockType.DBT,
             language=BlockLanguage.SQL,
             configuration={
-                'file_path': str(Path('test_project_name/models/model.sql')),
+                'file_path': os.path.join(project_path, 'models', 'demo', 'model.sql'),
                 'dbt': {'disable_tests': True}
             },
             pipeline=pipeline,
@@ -42,9 +56,7 @@ class DBTBlockTest(TestCase):
             block_type=BlockType.DBT,
             language=BlockLanguage.YAML,
             pipeline=pipeline,
-            configuration={
-                'dbt_project_name': 'test_project_name'
-            }
+            configuration={},
         )
 
         self.dbt_block_yaml.upstream_blocks = [self.block_sql]
@@ -71,43 +83,63 @@ class DBTBlockTest(TestCase):
         self.assertFalse(isinstance(self.dbt_block_yaml, DBTBlockSQL))
         self.assertTrue(isinstance(self.dbt_block_yaml, DBTBlockYAML))
 
-    def test_base_project_path(self):
-        self.assertTrue(
-            Path(self.dbt_block_sql.base_project_path).match('test_repo_path/dbt')
+    def test_paths(self):
+        self.assertEqual(
+            self.dbt_block_sql.project_path,
+            os.path.join(self.repo_path, 'dbt/dir1/dir2/dir3/fire'),
+        )
+        self.assertEqual(
+            self.dbt_block_sql.file_path,
+            os.path.join(self.repo_path, 'dbt/dir1/dir2/dir3/fire/models/demo/model.sql'),
+        )
+
+        self.dbt_block_yaml.configuration = dict(
+            dbt_profiles_file_path='dbt/dir1/dir2/dir3/fire/models/demo',
+        )
+        self.assertEqual(
+            self.dbt_block_yaml.project_path,
+            os.path.join(self.repo_path, 'dbt/dir1/dir2/dir3/fire'),
+        )
+        self.dbt_block_yaml.configuration = dict(
+            dbt_project_name='dbt/dir1/dir2/dir3/fire/models',
+        )
+        self.assertEqual(
+            self.dbt_block_yaml.project_path,
+            os.path.join(self.repo_path, 'dbt/dir1/dir2/dir3/fire'),
         )
 
     def test_dbt_configuration(self):
         self.assertEqual(self.dbt_block_sql._dbt_configuration, {'disable_tests': True})
         self.assertEqual(self.dbt_block_yaml._dbt_configuration, {})
 
-    def test_variables_json(self):
-        self.assertEqual(
-            self.dbt_block_sql._variables_json({
-                'key1': 1,
-                'key2': 'foo',
-                'key3': ['bar']
-            }),
-            '{"key1": 1, "key2": "foo", "key3": ["bar"]}'
-        )
-
-    @patch('mage_ai.data_preparation.models.block.dbt.block.Profiles')
-    @patch('mage_ai.data_preparation.models.block.dbt.block.Project')
-    @patch('mage_ai.data_preparation.models.block.dbt.block.Sources')
-    @patch('mage_ai.data_preparation.models.block.dbt.block.DBTAdapter')
-    def test_update_sources(self, DBTAdapter, Sources, Project, Profiles):
-        DBTAdapter.return_value.__enter__.return_value.credentials.schema = 'public'
-        DBTAdapter.return_value.__enter__.return_value.credentials.database = None
-        Sources.return_value.reset_pipeline = MagicMock()
-
-        DBTBlock.update_sources({
-            'test_block_sql': self.block_sql,
-            'test_dbt_block_sql': self.dbt_block_sql,
-            'test_dbt_block_yaml': self.dbt_block_yaml
-        })
-
-        Sources.return_value.reset_pipeline.assert_called_once_with(
-            project_name='test_project_name',
-            pipeline_uuid='test',
-            block_uuids=['test_block_sql'],
-            schema='public'
-        )
+    # def test_variables_json(self):
+    #     self.assertEqual(
+    #         self.dbt_block_sql._variables_json({
+    #             'key1': 1,
+    #             'key2': 'foo',
+    #             'key3': ['bar']
+    #         }),
+    #         '{"key1": 1, "key2": "foo", "key3": ["bar"]}'
+    #     )
+
+    # @patch('mage_ai.data_preparation.models.block.dbt.block.Profiles')
+    # @patch('mage_ai.data_preparation.models.block.dbt.block.Project')
+    # @patch('mage_ai.data_preparation.models.block.dbt.block.Sources')
+    # @patch('mage_ai.data_preparation.models.block.dbt.block.DBTAdapter')
+    # def test_update_sources(self, DBTAdapter, Sources, Project, Profiles):
+    #     DBTAdapter.return_value.__enter__.return_value.credentials.schema = 'public'
+    #     DBTAdapter.return_value.__enter__.return_value.credentials.database = None
+    #     Sources.return_value.reset_pipeline = MagicMock()
+
+    #     DBTBlock.update_sources({
+    #         'test_block_sql': self.block_sql,
+    #         'test_dbt_block_sql': self.dbt_block_sql,
+    #         'test_dbt_block_yaml': self.dbt_block_yaml
+    #     })
+
+    #     Sources.return_value.reset_pipeline.assert_called_once_with(
+    #         project_name='test_project_name',
+    #         pipeline_uuid='test',
+    #         block_uuids=['test_block_sql'],
+    #         schema='public'
+    #     )