<commit_before>from mapreduce import output_writers
from django.db import models
from djangae.test import TestCase
from unittest import skipIf

from django.db import connection
from django.core.files.base import ContentFile
from djangae.storage import CloudStorage, has_cloudstorage
from django.test.utils import override_settings
from djangae.contrib.processing.mapreduce import (
    map_files,
    map_queryset,
    map_entities,
    map_reduce_queryset,
    map_reduce_entities,
    get_pipeline_by_id,
)

from google.appengine.api import datastore


class TestModel(models.Model):
    class Meta:
        app_label = "mapreduce"

    is_true = models.BooleanField(default=False)
    text = models.CharField(null=True, max_length=50)


class Counter(models.Model):
    count = models.PositiveIntegerField(default=0)


def count(instance, counter_id):
    counter = Counter.objects.get(pk=counter_id)
    counter.count = models.F('count') + 1
    counter.save()


def count_contents(file_obj, counter_id):
    counter = Counter.objects.get(pk=counter_id)
    counter.count = models.F('count') + len(file_obj.read())
    counter.save()


def yield_letters(instance):
    if hasattr(instance, 'text'):
        text = instance.text
    else:
        text = instance.get('text', '')
    for letter in text:
        yield (letter, "")


def reduce_count(key, values):
    yield (key, len(values))


def delete(*args, **kwargs):
    TestModel.objects.all().delete()


class MapReduceEntityTests(TestCase):

    def setUp(self):
        for i in range(5):
            TestModel.objects.create(
                id=i+1,
                text="abcc"
            )

    def test_filters(self):
        pipeline = map_reduce_entities(
            TestModel._meta.db_table,
            connection.settings_dict["NAMESPACE"],
            yield_letters,
            reduce_count,
            output_writers.GoogleCloudStorageKeyValueOutputWriter,
            _output_writer_kwargs={
                'bucket_name': 'test-bucket'
            },
            _filters=[("text", "=", "abcc")]
        )
        self.process_task_queues()
        # Refetch the pipeline record
        pipeline = get_pipeline_by_id(pipeline.pipeline_id)
        self.assertTrue(pipeline.has_finalized)

    def test_mapreduce_over_entities(self):
        pipeline = map_reduce_entities(
            TestModel._meta.db_table,
            connection.settings_dict["NAMESPACE"],
            yield_letters,
            reduce_count,
            output_writers.GoogleCloudStorageKeyValueOutputWriter,
            _output_writer_kwargs={
                'bucket_name': 'test-bucket'
            }
        )
        self.process_task_queues()
        # Refetch the pipeline record
        pipeline = get_pipeline_by_id(pipeline.pipeline_id)
        self.assertTrue(pipeline.has_finalized)


class MapReduceQuerysetTests(TestCase):
<commit_msg>Update `MapReduceEntityTests.test_filters` to test the filtering.<commit_after>from mapreduce import output_writers
from django.db import models
from djangae.test import TestCase
from unittest import skipIf

from django.db import connection
from django.core.files.base import ContentFile
from djangae.storage import CloudStorage, has_cloudstorage
from django.test.utils import override_settings
from djangae.contrib.processing.mapreduce import (
    map_files,
    map_queryset,
    map_entities,
    map_reduce_queryset,
    map_reduce_entities,
    get_pipeline_by_id,
)

from google.appengine.api import datastore


class TestModel(models.Model):
    class Meta:
        app_label = "mapreduce"

    is_true = models.BooleanField(default=False)
    text = models.CharField(null=True, max_length=50)


class Counter(models.Model):
    count = models.PositiveIntegerField(default=0)


def count(instance, counter_id):
    counter = Counter.objects.get(pk=counter_id)
    counter.count = models.F('count') + 1
    counter.save()


def count_entity_to_default_counter(entity):
    """ Dirty hack to work around the fact that when using  `map_reduce_entities` we cannot pass
        any params to the map_func, and therefore this function only accepts a single arg of the
        `entity` and just assumes that there's only 1 Counter object!
    """
    counter = Counter.objects.get()
    counter.count = models.F('count') + 1
    counter.save()


def count_contents(file_obj, counter_id):
    counter = Counter.objects.get(pk=counter_id)
    counter.count = models.F('count') + len(file_obj.read())
    counter.save()


def yield_letters(instance):
    if hasattr(instance, 'text'):
        text = instance.text
    else:
        text = instance.get('text', '')
    for letter in text:
        yield (letter, "")


def reduce_count(key, values):
    yield (key, len(values))


def delete(*args, **kwargs):
    TestModel.objects.all().delete()


class MapReduceEntityTests(TestCase):

    def setUp(self):
        for i in range(5):
            TestModel.objects.create(
                id=i+1,
                text="abcc-%s" % i
            )

    def test_filters(self):
        """ Passing the `_filters` kwarg to to `map_reduce_entities` should allow only some
            entities to be processed.
        """
        counter = Counter.objects.create()
        pipeline = map_reduce_entities(
            TestModel._meta.db_table,
            connection.settings_dict["NAMESPACE"],
            count_entity_to_default_counter,
            reduce_count,  # This is a no-op because count_entity doesn't return anything
            output_writers.GoogleCloudStorageKeyValueOutputWriter,
            _output_writer_kwargs={
                'bucket_name': 'test-bucket'
            },
            _filters=[("text", "=", "abcc-3")]
        )
        self.process_task_queues()
        # Refetch the pipeline record
        pipeline = get_pipeline_by_id(pipeline.pipeline_id)
        self.assertTrue(pipeline.has_finalized)
        # We expect only the one entity to have been counted
        counter.refresh_from_db()
        self.assertEqual(counter.count, 1)

    def test_mapreduce_over_entities(self):
        pipeline = map_reduce_entities(
            TestModel._meta.db_table,
            connection.settings_dict["NAMESPACE"],
            yield_letters,
            reduce_count,
            output_writers.GoogleCloudStorageKeyValueOutputWriter,
            _output_writer_kwargs={
                'bucket_name': 'test-bucket'
            }
        )
        self.process_task_queues()
        # Refetch the pipeline record
        pipeline = get_pipeline_by_id(pipeline.pipeline_id)
        self.assertTrue(pipeline.has_finalized)


class MapReduceQuerysetTests(TestCase):

{'fraction_non_alphanumeric': 0.09069140975755761, 'fraction_numerical': 0.002394492666866208, 'mean_word_length': 2.807977207977208, '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': '414385', 'n_tokens_mistral': 1949, 'n_tokens_neox': 1930, 'n_words': 414}