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

mage_ai/tests/data_cleaner/transformer_actions/test_trim_transformer.py

@@ -0,0 +1,161 @@
+import unittest
+
+import pandas as pd
+
+from mage_ai.data_cleaner.transformer_actions.base import BaseAction
+from mage_ai.data_cleaner.transformer_actions.constants import ActionType, Axis
+from mage_ai.data_cleaner.transformer_actions.utils import build_transformer_action
+
+
+class TestTrimTransformer(unittest.TestCase):
+    def test_basic_trimming(self):
+        df = pd.DataFrame({
+            'category': ['  A  ', '  B  ', '  C  ', '  D  ']
+        })
+
+        expected_output = pd.DataFrame({
+            'category': ['A', 'B', 'C', 'D']
+        })
+
+        action = build_transformer_action(
+            df,
+            action_type=ActionType.REFORMAT,
+            arguments=['category'],
+            axis=Axis.COLUMN,
+            options={'reformat': 'trim'}
+        )
+
+        output = BaseAction(action).execute(df)
+
+        pd.testing.assert_frame_equal(output, expected_output)
+        print("Test 1 (Basic Trimming): Passed")
+
+    def test_no_trimming_needed(self):
+        df = pd.DataFrame({
+            'category': ['A', 'B', 'C', 'D']
+        })
+
+        expected_output = df.copy()
+
+        action = build_transformer_action(
+            df,
+            action_type=ActionType.REFORMAT,
+            arguments=['category'],
+            axis=Axis.COLUMN,
+            options={'reformat': 'trim'}
+        )
+
+        output = BaseAction(action).execute(df)
+
+        pd.testing.assert_frame_equal(output, expected_output)
+        print("Test 2 (No Trimming Needed): Passed")
+
+    def test_mixed_content(self):
+        df = pd.DataFrame({
+            'category': ['  A  ', 'B', '  C', 'D  ']
+        })
+
+        expected_output = pd.DataFrame({
+            'category': ['A', 'B', 'C', 'D']
+        })
+
+        action = build_transformer_action(
+            df,
+            action_type=ActionType.REFORMAT,
+            arguments=['category'],
+            axis=Axis.COLUMN,
+            options={'reformat': 'trim'}
+        )
+
+        output = BaseAction(action).execute(df)
+
+        pd.testing.assert_frame_equal(output, expected_output)
+        print("Test 3 (Mixed Content): Passed")
+
+    def test_multiple_columns(self):
+        df = pd.DataFrame({
+            'category': ['  A  ', '  B  ', '  C  ', '  D  '],
+            'subcategory': ['  X  ', 'Y  ', '  Z', ' W ']
+        })
+
+        expected_output = pd.DataFrame({
+            'category': ['A', 'B', 'C', 'D'],
+            'subcategory': ['X', 'Y', 'Z', 'W']
+        })
+
+        action = build_transformer_action(
+            df,
+            action_type=ActionType.REFORMAT,
+            arguments=['category', 'subcategory'],
+            axis=Axis.COLUMN,
+            options={'reformat': 'trim'}
+        )
+
+        output = BaseAction(action).execute(df)
+
+        pd.testing.assert_frame_equal(output, expected_output)
+        print("Test 4 (Multiple Columns): Passed")
+
+    def test_empty_strings(self):
+        df = pd.DataFrame({
+            'category': ['  ', '  B  ', '  ', '  D  ']
+        })
+
+        expected_output = pd.DataFrame({
+            'category': ['', 'B', '', 'D']
+        })
+
+        action = build_transformer_action(
+            df,
+            action_type=ActionType.REFORMAT,
+            arguments=['category'],
+            axis=Axis.COLUMN,
+            options={'reformat': 'trim'}
+        )
+
+        output = BaseAction(action).execute(df)
+
+        pd.testing.assert_frame_equal(output, expected_output)
+        print("Test 5 (Empty Strings): Passed")
+
+    def test_no_columns_to_trim(self):
+        df = pd.DataFrame({
+            'category': ['  A  ', '  B  ', '  C  ', '  D  ']
+        })
+
+        expected_output = df.copy()
+
+        action = build_transformer_action(
+            df,
+            action_type=ActionType.REFORMAT,
+            arguments=[],
+            axis=Axis.COLUMN,
+            options={'reformat': 'trim'}
+        )
+
+        output = BaseAction(action).execute(df)
+
+        pd.testing.assert_frame_equal(output, expected_output)
+        print("Test 6 (No Columns to Trim): Passed")
+
+    def test_column_not_found(self):
+        df = pd.DataFrame({
+            'category': ['  A  ', '  B  ', '  C  ', '  D  ']
+        })
+
+        with self.assertRaises(KeyError):
+            action = build_transformer_action(
+                df,
+                action_type=ActionType.REFORMAT,
+                arguments=['nonexistent_column'],
+                axis=Axis.COLUMN,
+                options={'reformat': 'trim'}
+            )
+
+            BaseAction(action).execute(df)
+
+        print("Test 7 (Column Not Found): Passed")
+
+
+if __name__ == '__main__':
+    unittest.main()

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

@@ -85,7 +85,7 @@ target: dev
             lambda: MockUUID(),
         ):
             with Profiles(self.repo_path, self.variables) as profiles:
-                self.assertNotEquals(
+                self.assertNotEqual(
                     profiles.profiles_dir,
                     self.repo_path
                 )

mage_ai/tests/data_preparation/models/block/dynamic/test_combos.py

@@ -110,7 +110,7 @@ class DynamicBlockCombinationTest(BaseApiTestCase):
             content = inspect.getsource(func)
 
         if content:
-            block.update_content(f'@{block.type}\n{content}')
+            block.update_content(block.get_typed_content(content))
 
         return block
 

mage_ai/tests/data_preparation/models/block/hook/test_hook_block.py

@@ -5,6 +5,7 @@ from faker import Faker
 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.global_hooks.models import HookStatus, HookStrategy
+from mage_ai.shared.constants import ENV_TEST
 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
@@ -45,7 +46,7 @@ class HookBlockTest(BaseApiTestCase):
                 check_status=True,
                 configuration={},
                 context={},
-                env='test',
+                env=ENV_TEST,
                 error_on_failure=True,
                 mage=1,
                 pipeline_uuid='',
@@ -84,7 +85,7 @@ class HookBlockTest(BaseApiTestCase):
                 check_status=True,
                 configuration={},
                 context={},
-                env='test',
+                env=ENV_TEST,
                 error_on_failure=True,
                 mage=1,
                 pipeline_uuid='',

mage_ai/tests/data_preparation/models/test_blocks_helper.py

@@ -107,7 +107,7 @@ class BlockHelperTest(BaseApiTestCase):
                 '\n'.join([
                     'from mage_ai.data.models.generator import DataGenerator',
                     'from mage_ai.data.tabular.mocks import create_dataframe',
-                    f'@{block.type}\n{content}',
+                    block.get_typed_content(content),
                 ])
             )
 

mage_ai/tests/data_preparation/models/variables/test_summarizer.py

@@ -110,7 +110,7 @@ class VariableSummarizerTest(BaseApiTestCase):
                 '\n'.join([
                     'from mage_ai.data.models.generator import DataGenerator',
                     'from mage_ai.data.tabular.mocks import create_dataframe',
-                    f'@{block.type}\n{content}',
+                    block.get_typed_content(content),
                 ])
             )
 

mage_ai/tests/data_preparation/sync/test_git_sync.py

@@ -22,8 +22,8 @@ class GitSyncTest(DBTestCase):
 
         git_sync.sync_data()
 
-        git_manager_instance_mock.reset_hard.called_once_with(branch='dev')
-        git_manager_instance_mock.submodules_update.not_called()
+        git_manager_instance_mock.reset_hard.assert_called_once_with(branch='dev')
+        git_manager_instance_mock.submodules_update.assert_not_called()
 
         with open(os.path.join(self.repo_path, '.preferences.yaml'), 'r') as f:
             preferences = yaml.safe_load(f)
@@ -42,8 +42,8 @@ class GitSyncTest(DBTestCase):
 
         git_sync.sync_data()
 
-        git_manager_instance_mock.reset_hard.called_once_with(branch='dev')
-        git_manager_instance_mock.submodules_update.called_once()
+        git_manager_instance_mock.reset_hard.assert_called_once_with(branch='dev')
+        git_manager_instance_mock.submodules_update.assert_called_once()
 
         with open(os.path.join(self.repo_path, '.preferences.yaml'), 'r') as f:
             preferences = yaml.safe_load(f)
@@ -63,7 +63,7 @@ class GitSyncTest(DBTestCase):
 
         git_sync.sync_data()
 
-        git_manager_instance_mock.reset_hard.called_once_with(branch='dev')
+        git_manager_instance_mock.reset_hard.assert_called_once_with(branch='dev')
 
         self.assertFalse(os.path.exists(preferences_file_path))
 
@@ -77,4 +77,4 @@ class GitSyncTest(DBTestCase):
 
         git_sync.reset()
 
-        git_manager_instance_mock.clone.called_once_with(sync_submodules=False)
+        git_manager_instance_mock.clone.assert_called_once_with(sync_submodules=False)

mage_ai/tests/orchestration/queue/test_process_queue.py

@@ -33,9 +33,10 @@ class ProcessQueueTests(TestCase):
         self.assertFalse('block_run_3' in self.queue.job_dict)
         self.assertFalse('block_run_4' in self.queue.job_dict)
 
-    @patch('mage_ai.orchestration.queue.process_queue.poll_job_and_execute')
+    @patch.object(ProcessQueue, 'start_worker_pool')
     @patch('mage_ai.orchestration.queue.process_queue.psutil.pid_exists')
-    def test_has_job(self, mock_pid_exists, mock_poll_job_and_execute):
+    def test_has_job(self, mock_pid_exists, mock_start_worker_pool):
+        mock_start_worker_pool.return_value = None
         mock_pid_exists.return_value = True
 
         self.queue.job_dict['block_run_1'] = JobStatus.QUEUED

mage_ai/tests/orchestration/test_pipeline_scheduler.py

@@ -441,6 +441,7 @@ class PipelineSchedulerTests(DBTestCase):
                 execution_partition=f'{pipeline_run.pipeline_schedule_id}',
                 event={},
                 pipeline_run_id=pipeline_run.id,
+                trigger_name='test_pipeline_2_trigger'
             ),
             dict(
                 pipeline_run_id=pipeline_run.id,

mage_ai/tests/settings/test_platform.py

@@ -205,12 +205,12 @@ class PlatformSettingsTest(ProjectPlatformMixin):
         ))
 
     def test_platform_settings_full_path(self):
-        self.assertEquals(
+        self.assertEqual(
             platform_settings_full_path(), os.path.join(base_repo_path(), 'settings.yaml'),
         )
 
     def test_local_platform_settings_full_path(self):
-        self.assertEquals(
+        self.assertEqual(
             local_platform_settings_full_path(), os.path.join(base_repo_path(), '.settings.yaml'),
         )
 

mage_ai/tests/streaming/sinks/test_generic_io.py

@@ -1,4 +1,4 @@
-import os
+from pathlib import Path
 from unittest.mock import ANY, MagicMock, call, patch
 
 import pandas as pd
@@ -52,33 +52,40 @@ TEST_DATABASES = [
 
 
 class GenericIOTests(TestCase):
-    @patch('mage_ai.streaming.sinks.generic_io.ConfigFileLoader')
+    def setUp(self):
+        super().setUp()
+        self.test_path = Path(self.repo_path)
+        self.test_config_path = self.test_path / 'io_config.yaml'
+        sample_yaml = """
+test_profile:
+  test_1: test_1
+"""
+        with self.test_config_path.open('w') as fout:
+            fout.write(sample_yaml)
+
+    def tearDown(self):
+        self.test_config_path.unlink()
+        super().tearDown()
+
     @patch('mage_ai.streaming.sinks.generic_io.importlib.import_module')
-    def test_init_client(self, mock_import_module, mock_config_loader):
-        first_db = True
+    def test_init_client(self, mock_import_module):
         for database in TEST_DATABASES:
-            mock_objects = self.__mock_objects(mock_import_module, mock_config_loader, database)
+            mock_objects = self.__mock_objects(mock_import_module, database)
 
             # Assertions
-            if first_db:
-                mock_config_loader.assert_called_once_with(
-                    os.path.join(self.repo_path, 'io_config.yaml'), 'test_profile')
-                mock_import_module.assert_called_once_with(database['module_path'])
-            mock_objects['io_class'].with_config.assert_called_once_with(
-                mock_objects['config_loader_instance'])
+            mock_import_module.assert_any_call(database['module_path'])
+            mock_objects['io_class'].with_config.assert_called_once()
             mock_objects['io_client'].open.assert_called_once()
 
             self.assertEqual(mock_objects['sink'].io_client, mock_objects['io_client'])
 
             del mock_objects['sink']
             mock_objects['io_client'].close.assert_called_once()
-            first_db = False
 
-    @patch('mage_ai.streaming.sinks.generic_io.ConfigFileLoader')
     @patch('mage_ai.streaming.sinks.generic_io.importlib.import_module')
-    def test_write(self, mock_import_module, mock_config_loader):
+    def test_write(self, mock_import_module):
         database = TEST_DATABASES[0]
-        mock_objects = self.__mock_objects(mock_import_module, mock_config_loader, database)
+        mock_objects = self.__mock_objects(mock_import_module, database)
 
         mock_objects['sink'].batch_write = MagicMock()
 
@@ -90,11 +97,10 @@ class GenericIOTests(TestCase):
         # Assertions
         mock_objects['sink'].batch_write.assert_called_once_with([message])
 
-    @patch('mage_ai.streaming.sinks.generic_io.ConfigFileLoader')
     @patch('mage_ai.streaming.sinks.generic_io.importlib.import_module')
-    def test_batch_write(self, mock_import_module, mock_config_loader):
+    def test_batch_write(self, mock_import_module):
         database = TEST_DATABASES[0]
-        mock_objects = self.__mock_objects(mock_import_module, mock_config_loader, database)
+        mock_objects = self.__mock_objects(mock_import_module, database)
         mock_objects['io_client'].export = MagicMock()
 
         # Test data
@@ -130,13 +136,12 @@ class GenericIOTests(TestCase):
             ]),
         )
 
-    def __mock_objects(self, mock_import_module, mock_config_loader, database):
+    def __mock_objects(self, mock_import_module, database):
         mock_io_module = MagicMock()
         mock_io_class = MagicMock()
         mock_io_client = MagicMock()
         mock_import_module.return_value = mock_io_module
         mock_io_class.with_config.return_value = mock_io_client
-        mock_config_loader_instance = mock_config_loader.return_value
         setattr(mock_io_module, database['class_name'], mock_io_class)
         generic_io_sink = GenericIOSink(dict(
             connector_type=database['connector_type'],
@@ -146,7 +151,6 @@ class GenericIOTests(TestCase):
             )
         ))
         return dict(
-            config_loader_instance=mock_config_loader_instance,
             io_module=mock_io_module,
             io_class=mock_io_class,
             io_client=mock_io_client,

mage_ai/usage_statistics/constants.py

@@ -1,9 +1,9 @@
-import enum
+from mage_ai.shared.enum import StrEnum
 
 API_ENDPOINT = 'https://api.mage.ai/v1/usage_statistics'
 
 
-class EventNameType(str, enum.Enum):
+class EventNameType(StrEnum):
     API_ERROR = 'api_error'
     APPLICATION_ERROR = 'application_error'
     BLOCK_RUN_ENDED = 'block_run_ended'
@@ -13,14 +13,14 @@ class EventNameType(str, enum.Enum):
     USAGE_STATISTIC_CREATE = 'usage_statistic.create'
 
 
-class EventActionType(str, enum.Enum):
+class EventActionType(StrEnum):
     CREATE = 'create'
     DENY = 'deny'
     EXECUTE = 'execute'
     IMPRESSION = 'impression'
 
 
-class EventObjectType(str, enum.Enum):
+class EventObjectType(StrEnum):
     BLOCK = 'block'
     BLOCK_RUN = 'block_run'
     CHART = 'chart'