mage-ai 0.9.32__py3-none-any.whl → 0.9.34__py3-none-any.whl

mage_ai/tests/data_preparation/models/block/sql/utils/test_shared.py

@@ -1,14 +1,141 @@
 from mage_ai.data_preparation.models.block import Block
-from mage_ai.data_preparation.models.block.sql.utils.shared import interpolate_input
+from mage_ai.data_preparation.models.block.sql.utils.shared import (
+    extract_create_statement_table_name,
+    extract_drop_statement_table_names,
+    extract_insert_statement_table_names,
+    extract_update_statement_table_names,
+    interpolate_input,
+)
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.tests.base_test import DBTestCase
 
 
 class SQLBlockSharedUtilsTest(DBTestCase):
-    def setUp(self):
-        super().setUp()
+    def test_extract_create_statement_table_name(self):
+        # Test case 1: Basic CREATE TABLE statement
+        text = """
+        -- This is a comment
+        CREATE TABLE table_name (
+            column1 INT,
+            column2 VARCHAR(255)
+        );
+        """
+        expected_result = 'table_name'
+        self.assertEqual(extract_create_statement_table_name(text), expected_result)
+
+        # Test case 2: CREATE TABLE IF NOT EXISTS
+        text = """
+        CREATE TABLE IF NOT EXISTS schema.table_name (
+            column1 INT,
+            column2 VARCHAR(255)
+        );
+        """
+        expected_result = 'schema.table_name'
+        self.assertEqual(extract_create_statement_table_name(text), expected_result)
+
+        # Test case 3: No CREATE TABLE statement
+        text = """
+        SELECT * FROM table_name;
+        """
+        self.assertIsNone(extract_create_statement_table_name(text))
+
+        # Test case 4: CREATE TABLE with additional clauses
+        text = """
+        CREATE TABLE IF NOT EXISTS schema.table_name
+        ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'
+        WITH SERDEPROPERTIES (
+            "separatorChar" = "\\t",
+            "quoteChar"     = "'",
+            "escapeChar"    = "\\\\"
+        )
+        STORED AS TEXTFILE;
+        """
+        expected_result = 'schema.table_name'
+        self.assertEqual(extract_create_statement_table_name(text), expected_result)
+
+    def test_extract_insert_statement_table_names(self):
+        # Test case 1: Basic test with one INSERT statement
+        text = """
+        -- This is a comment
+        INSERT INTO table_name
+        VALUES (1, 'John');
+        """
+        expected_result = ['table_name']
+        self.assertEqual(extract_insert_statement_table_names(text), expected_result)
+
+        # Test case 2: INSERT INTO with overwrite
+        text = """
+        INSERT OVERWRITE INTO database.table1
+        SELECT * FROM source_table;
+        """
+        expected_result = ['database.table1']
+        self.assertEqual(extract_insert_statement_table_names(text), expected_result)
+
+        # Test case 3: No INSERT statements
+        text = """
+        SELECT * FROM table_name;
+        """
+        expected_result = []
+        self.assertEqual(extract_insert_statement_table_names(text), expected_result)
+
+    def test_extract_drop_statement_table_names(self):
+        # Test case 1: Basic DROP TABLE statement
+        text = """
+        -- This is a comment
+        DROP TABLE table_name;
+        """
+        expected_result = ['table_name']
+        self.assertEqual(extract_drop_statement_table_names(text), expected_result)
+
+        # Test case 2: DROP TABLE IF EXISTS
+        text = """
+        drop table if exists schema.table1;
+        """
+        expected_result = ['schema.table1']
+        self.assertEqual(extract_drop_statement_table_names(text), expected_result)
 
+        # Test case 3: No DROP TABLE statements
+        text = """
+        SELECT * FROM table_name;
+        """
+        expected_result = []
+        self.assertEqual(extract_drop_statement_table_names(text), expected_result)
+
+    def test_extract_update_statement_table_names(self):
+        # Test case 1: Basic test with one UPDATE statement
+        text = """
+        -- This is a comment
+        UPDATE table_name
+        SET column1 = value1, column2 = value2
+        WHERE condition;
+        """
+        expected_result = ['table_name']
+        self.assertEqual(extract_update_statement_table_names(text), expected_result)
+
+        # Test case 2: Multiple UPDATE statements
+        text = """
+        UPDATE table1
+        SET column1 = value1, column2 = value2
+        WHERE condition;
+
+        Another UPDATE statement here.
+
+        update table2
+        set column3 = value3
+        where another_condition;
+        """
+        expected_result = ['table1', 'table2']
+        self.assertEqual(extract_update_statement_table_names(text), expected_result)
+
+        # Test case 3: No UPDATE statements
+        text = """
+        SELECT * FROM table_name;
+        """
+        expected_result = []
+        self.assertEqual(extract_update_statement_table_names(text), expected_result)
+
+    def test_interpolate_input(self):
         self.pipeline = Pipeline.create(
             'test pipeline',
             repo_path=self.repo_path,
@@ -102,8 +229,6 @@ ON 1 = 1
             data_loader_1.uuid,
             data_loader_2.uuid,
         ])
-
-    def test_interpolate_input(self):
         data_exporter = self.pipeline.get_block('data_exporter')
 
         query_string = interpolate_input(

mage_ai/tests/factory.py

@@ -6,6 +6,7 @@ from faker import Faker
 from mage_ai.authentication.passwords import create_bcrypt_hash, generate_salt
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.models.triggers import ScheduleType
 from mage_ai.orchestration.db.models.oauth import User
 from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
 from mage_ai.shared.hash import merge_dict
@@ -87,6 +88,7 @@ def create_pipeline_run_with_schedule(
         pipeline_schedule = PipelineSchedule.create(
             name=f'{pipeline_uuid}_trigger',
             pipeline_uuid=pipeline_uuid,
+            schedule_type=ScheduleType.TIME,
             settings=pipeline_schedule_settings,
         )
         pipeline_schedule_id = pipeline_schedule.id

mage_ai/tests/orchestration/test_pipeline_scheduler.py

@@ -26,7 +26,7 @@ from mage_ai.orchestration.pipeline_scheduler import (
     schedule_all,
 )
 from mage_ai.shared.array import find
-from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.hash import ignore_keys, merge_dict
 from mage_ai.tests.base_test import DBTestCase
 from mage_ai.tests.factory import (
     create_pipeline_run_with_schedule,
@@ -110,6 +110,7 @@ class PipelineSchedulerTests(DBTestCase):
             print(f'Mock PipelineScheduler().schedule() for pipeline run {self.pipeline_run.id}.')
 
         shared_attrs = dict(
+            name='trigger',
             pipeline_uuid='test_pipeline',
             schedule_interval=ScheduleInterval.HOURLY,
             schedule_type=ScheduleType.TIME,
@@ -151,6 +152,7 @@ class PipelineSchedulerTests(DBTestCase):
             print(f'Mock PipelineScheduler().schedule() for pipeline run {self.pipeline_run.id}.')
 
         shared_attrs = dict(
+            name='trigger',
             pipeline_uuid='test_pipeline',
             schedule_interval=ScheduleInterval.HOURLY,
             schedule_type=ScheduleType.TIME,
@@ -272,7 +274,8 @@ class PipelineSchedulerTests(DBTestCase):
         mock_has_pipeline_run_job.return_value = False
         mock_job_manager.add_job = MagicMock()
         scheduler.schedule()
-        mock_job_manager.add_job.assert_called_once_with(
+
+        call_args = [
             JobType.PIPELINE_RUN,
             pipeline_run.id,
             mock_run_pipeline,
@@ -288,7 +291,16 @@ class PipelineSchedulerTests(DBTestCase):
                 pipeline_schedule_id=pipeline_run.pipeline_schedule_id,
                 pipeline_uuid='test_pipeline_2',
             ),
-        )
+        ]
+
+        self.assertTrue(len(mock_job_manager.add_job.mock_calls) == 1)
+
+        mock_call = mock_job_manager.add_job.mock_calls[0]
+
+        for i in range(5):
+            self.assertEqual(call_args[i], mock_call[1][i])
+
+        self.assertEqual(call_args[5], ignore_keys(mock_call[1][5], ['hostname']))
 
     def test_on_block_complete(self):
         pipeline_run = create_pipeline_run_with_schedule(pipeline_uuid='test_pipeline')
@@ -458,6 +470,7 @@ class PipelineSchedulerTests(DBTestCase):
         pipeline_schedule = PipelineSchedule.create(
             name='test_sla_pipeline_trigger',
             pipeline_uuid=pipeline_uuid,
+            schedule_type=ScheduleType.TIME,
             sla=600,
         )
         pipeline_schedule.update(
@@ -501,6 +514,7 @@ class PipelineSchedulerTests(DBTestCase):
         PipelineSchedule.create(
             name='test_sla_pipeline_trigger_1',
             pipeline_uuid=pipeline.uuid,
+            schedule_type=ScheduleType.TIME,
             status=ScheduleStatus.ACTIVE,
             start_time=datetime(2023, 4, 1, 1, 20, 33),
             schedule_interval='@hourly',
@@ -508,6 +522,7 @@ class PipelineSchedulerTests(DBTestCase):
         PipelineSchedule.create(
             name='test_sla_pipeline_trigger_2',
             pipeline_uuid=pipeline.uuid,
+            schedule_type=ScheduleType.TIME,
             status=ScheduleStatus.ACTIVE,
             start_time=datetime(2023, 4, 5, 1, 20, 33),
             schedule_interval='@hourly',
@@ -536,6 +551,7 @@ class PipelineSchedulerTests(DBTestCase):
         pipeline_schedule = PipelineSchedule.create(
             name='test_timeout_pipeline_trigger',
             pipeline_uuid=pipeline_uuid,
+            schedule_type=ScheduleType.TIME,
             settings=dict(timeout=600),
         )
         pipeline_schedule.update(
@@ -583,6 +599,7 @@ class PipelineSchedulerTests(DBTestCase):
         pipeline_schedule = PipelineSchedule.create(
             name='test_block_timeout_pipeline_trigger',
             pipeline_uuid=pipeline_uuid,
+            schedule_type=ScheduleType.TIME,
         )
 
         block = pipeline.get_block('block1')

mage_ai/tests/streaming/sinks/test_google_cloud_pubsub.py

@@ -0,0 +1,28 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sinks.google_cloud_pubsub import GoogleCloudPubSubSink
+from mage_ai.tests.base_test import TestCase
+
+
+class GoogleCloudPubSubTests(TestCase):
+    def test_init(self):
+        with patch.object(GoogleCloudPubSubSink, 'init_client') as mock_init_client:
+            GoogleCloudPubSubSink(dict(
+                connector_type='test_google_cloud_pubsub',
+                project_id='test_project_id',
+                topic_id='test_topic_id',
+            ))
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        with patch.object(GoogleCloudPubSubSink, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                GoogleCloudPubSubSink(dict(
+                    connector_type='test_google_cloud_pubsub',
+                    project_id='test_project_id',
+                ))
+            self.assertTrue(
+                '__init__() missing 1 required positional argument: \'topic_id\''
+                in str(context.exception),
+            )
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/usage_statistics/logger.py

@@ -25,7 +25,7 @@ from mage_ai.data_preparation.models.project import Project
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import User
 from mage_ai.orchestration.db.models.schedules import PipelineRun
-from mage_ai.shared.environments import get_env
+from mage_ai.shared.environments import get_env, is_test
 from mage_ai.shared.hash import merge_dict
 from mage_ai.usage_statistics.constants import (
     API_ENDPOINT,
@@ -290,6 +290,9 @@ class UsageStatisticLogger():
         data: Dict,
         event_name: EventNameType = EventNameType.USAGE_STATISTIC_CREATE,
     ) -> bool:
+        if is_test():
+            return False
+
         if not self.help_improve_mage:
             return False
 

{mage_ai-0.9.32.dist-info → mage_ai-0.9.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.32
+Version: 0.9.34
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage