def test_convert_type(self, value, schema_type, expected):
     op = MySqlToGoogleCloudStorageOperator(task_id=TASK_ID,
                                            mysql_conn_id=MYSQL_CONN_ID,
                                            sql=SQL,
                                            bucket=BUCKET,
                                            filename=JSON_FILENAME)
     self.assertEqual(op.convert_type(value, schema_type), expected)
    def test_exec_success_json(self, gcs_hook_mock_class,
                               mysql_hook_mock_class):
        """Test successful run of execute function for JSON"""
        op = MySqlToGoogleCloudStorageOperator(task_id=TASK_ID,
                                               mysql_conn_id=MYSQL_CONN_ID,
                                               sql=SQL,
                                               bucket=BUCKET,
                                               filename=JSON_FILENAME)

        mysql_hook_mock = mysql_hook_mock_class.return_value
        mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
        mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION

        gcs_hook_mock = gcs_hook_mock_class.return_value

        def _assert_upload(bucket,
                           obj,
                           tmp_filename,
                           mime_type=None,
                           gzip=False):
            self.assertEqual(BUCKET, bucket)
            self.assertEqual(JSON_FILENAME.format(0), obj)
            self.assertEqual('application/json', mime_type)
            self.assertFalse(gzip)
            with open(tmp_filename, 'rb') as file:
                self.assertEqual(b''.join(NDJSON_LINES), file.read())

        gcs_hook_mock.upload.side_effect = _assert_upload

        op.execute(None)

        mysql_hook_mock_class.assert_called_once_with(
            mysql_conn_id=MYSQL_CONN_ID)
        mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(
            SQL)
    def test_schema_file(self, gcs_hook_mock_class, mysql_hook_mock_class):
        """Test writing schema files."""
        mysql_hook_mock = mysql_hook_mock_class.return_value
        mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
        mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION

        gcs_hook_mock = gcs_hook_mock_class.return_value

        def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip):  # pylint: disable=unused-argument
            if obj == SCHEMA_FILENAME:
                self.assertFalse(gzip)
                with open(tmp_filename, 'rb') as file:
                    self.assertEqual(b''.join(SCHEMA_JSON), file.read())

        gcs_hook_mock.upload.side_effect = _assert_upload

        op = MySqlToGoogleCloudStorageOperator(task_id=TASK_ID,
                                               sql=SQL,
                                               bucket=BUCKET,
                                               filename=JSON_FILENAME,
                                               schema_filename=SCHEMA_FILENAME)
        op.execute(None)

        # once for the file and once for the schema
        self.assertEqual(2, gcs_hook_mock.upload.call_count)
    def test_file_splitting(self, gcs_hook_mock_class, mysql_hook_mock_class):
        """Test that ndjson is split by approx_max_file_size_bytes param."""
        mysql_hook_mock = mysql_hook_mock_class.return_value
        mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
        mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION

        gcs_hook_mock = gcs_hook_mock_class.return_value
        expected_upload = {
            JSON_FILENAME.format(0): b''.join(NDJSON_LINES[:2]),
            JSON_FILENAME.format(1): NDJSON_LINES[2],
        }

        def _assert_upload(bucket,
                           obj,
                           tmp_filename,
                           mime_type=None,
                           gzip=False):
            self.assertEqual(BUCKET, bucket)
            self.assertEqual('application/json', mime_type)
            self.assertFalse(gzip)
            with open(tmp_filename, 'rb') as file:
                self.assertEqual(expected_upload[obj], file.read())

        gcs_hook_mock.upload.side_effect = _assert_upload

        op = MySqlToGoogleCloudStorageOperator(
            task_id=TASK_ID,
            sql=SQL,
            bucket=BUCKET,
            filename=JSON_FILENAME,
            approx_max_file_size_bytes=len(
                expected_upload[JSON_FILENAME.format(0)]))
        op.execute(None)
    def test_exec_success_csv_with_delimiter(self, gcs_hook_mock_class, mysql_hook_mock_class):
        """Test successful run of execute function for CSV with a field delimiter"""
        op = MySqlToGoogleCloudStorageOperator(
            task_id=TASK_ID,
            mysql_conn_id=MYSQL_CONN_ID,
            sql=SQL,
            export_format='csv',
            field_delimiter='|',
            bucket=BUCKET,
            filename=CSV_FILENAME)

        mysql_hook_mock = mysql_hook_mock_class.return_value
        mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
        mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION

        gcs_hook_mock = gcs_hook_mock_class.return_value

        def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False):
            self.assertEqual(BUCKET, bucket)
            self.assertEqual(CSV_FILENAME.format(0), obj)
            self.assertEqual('text/csv', mime_type)
            self.assertFalse(gzip)
            with open(tmp_filename, 'rb') as file:
                self.assertEqual(b''.join(CSV_LINES_PIPE_DELIMITED), file.read())

        gcs_hook_mock.upload.side_effect = _assert_upload

        op.execute(None)

        mysql_hook_mock_class.assert_called_once_with(mysql_conn_id=MYSQL_CONN_ID)
        mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL)
 def test_execute_with_query_error(self, mock_gcs_hook, mock_mysql_hook):
     mock_mysql_hook.return_value.get_conn.\
         return_value.cursor.return_value.execute.side_effect = ProgrammingError
     op = MySqlToGoogleCloudStorageOperator(task_id=TASK_ID,
                                            sql=SQL,
                                            bucket=BUCKET,
                                            filename=JSON_FILENAME,
                                            schema_filename=SCHEMA_FILENAME)
     with self.assertRaises(ProgrammingError):
         op.execute(None)
 def test_init(self):
     """Test MySqlToGoogleCloudStorageOperator instance is properly initialized."""
     op = MySqlToGoogleCloudStorageOperator(
         task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME,
         export_format='CSV', field_delimiter='|')
     self.assertEqual(op.task_id, TASK_ID)
     self.assertEqual(op.sql, SQL)
     self.assertEqual(op.bucket, BUCKET)
     self.assertEqual(op.filename, JSON_FILENAME)
     self.assertEqual(op.export_format, 'csv')
     self.assertEqual(op.field_delimiter, '|')