<commit_before>import os
import unittest
from unittest import mock
from tempfile import TemporaryDirectory
from contextlib import contextmanager
from cumulusci.tests.util import mock_salesforce_client, mock_describe_calls

from sqlalchemy import create_engine
from sqlalchemy.orm import create_session

from cumulusci.core.exceptions import TaskOptionsError, BulkDataException
from cumulusci.tasks.bulkdata import ExtractData
from cumulusci.tasks.bulkdata.step import (
    BaseQueryOperation,
    DataOperationStatus,
    DataOperationJobResult,
    DataOperationType,
    DataApi,
)
from cumulusci.tasks.bulkdata.tests.utils import _make_task
from cumulusci.tests.util import assert_max_memory_usage
from cumulusci.utils import temporary_dir
from cumulusci.tasks.bulkdata.mapping_parser import MappingLookup, MappingStep

import responses


@contextmanager
def mock_extract_jobs(task, extracted_records):
    def _job_state_from_batches(self, job_id):
        return DataOperationJobResult(
            DataOperationStatus.SUCCESS,
            [],
            10,
            0,
        )

    def get_results(self):
        return extracted_records[self.sobject]

    with mock.patch(
        "cumulusci.tasks.bulkdata.step.BulkApiQueryOperation.get_results",
        get_results,
    ), mock.patch(
        "cumulusci.tasks.bulkdata.step.BulkJobMixin._job_state_from_batches",
        _job_state_from_batches,
    ), mock.patch(
        "cumulusci.tasks.bulkdata.step.BulkApiQueryOperation.get_results",
        get_results,
    ):
        yield


class MockBulkQueryOperation(BaseQueryOperation):
    def __init__(self, *, sobject, api_options, context, query):
        super().__init__(
            sobject=sobject, api_options=api_options, context=context, query=query
        )
        self.results = []

    def query(self):
        self.job_id = "JOB"
        self.job_result = DataOperationJobResult(
            DataOperationStatus.SUCCESS, [], len(self.results), 0
        )

    def get_results(self):
        return iter(self.results)


class MockScalableBulkQueryOperation(MockBulkQueryOperation):
    def query(self):
        self.job_id = "JOB"
        self.job_result = DataOperationJobResult(
            DataOperationStatus.SUCCESS, [], self.result_len, 0
        )
<commit_msg>Remove double mock<commit_after>import os
import unittest
from unittest import mock
from tempfile import TemporaryDirectory
from contextlib import contextmanager
from cumulusci.tests.util import mock_salesforce_client, mock_describe_calls

from sqlalchemy import create_engine
from sqlalchemy.orm import create_session

from cumulusci.core.exceptions import TaskOptionsError, BulkDataException
from cumulusci.tasks.bulkdata import ExtractData
from cumulusci.tasks.bulkdata.step import (
    BaseQueryOperation,
    DataOperationStatus,
    DataOperationJobResult,
    DataOperationType,
    DataApi,
)
from cumulusci.tasks.bulkdata.tests.utils import _make_task
from cumulusci.tests.util import assert_max_memory_usage
from cumulusci.utils import temporary_dir
from cumulusci.tasks.bulkdata.mapping_parser import MappingLookup, MappingStep

import responses


@contextmanager
def mock_extract_jobs(task, extracted_records):
    def _job_state_from_batches(self, job_id):
        return DataOperationJobResult(
            DataOperationStatus.SUCCESS,
            [],
            10,
            0,
        )

    def get_results(self):
        return extracted_records[self.sobject]

    with mock.patch(
        "cumulusci.tasks.bulkdata.step.BulkApiQueryOperation.get_results",
        get_results,
    ), mock.patch(
        "cumulusci.tasks.bulkdata.step.BulkJobMixin._job_state_from_batches",
        _job_state_from_batches,
    ):
        yield


class MockBulkQueryOperation(BaseQueryOperation):
    def __init__(self, *, sobject, api_options, context, query):
        super().__init__(
            sobject=sobject, api_options=api_options, context=context, query=query
        )
        self.results = []

    def query(self):
        self.job_id = "JOB"
        self.job_result = DataOperationJobResult(
            DataOperationStatus.SUCCESS, [], len(self.results), 0
        )

    def get_results(self):
        return iter(self.results)


class MockScalableBulkQueryOperation(MockBulkQueryOperation):
    def query(self):
        self.job_id = "JOB"
        self.job_result = DataOperationJobResult(
            DataOperationStatus.SUCCESS, [], self.result_len, 0
        )

{'fraction_non_alphanumeric': 0.07330659536541889, 'fraction_numerical': 0.0022281639928698753, 'mean_word_length': 3.457795431976167, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2758447', 'n_tokens_mistral': 1332, 'n_tokens_neox': 1298, 'n_words': 259}