Ejemplo n.º 1
0
class ElasticsearchTestCase(TestCaseWithConfig):
    """Base class for Elastic Search related unit tests"""
    def setup_method(self, method):
        super().setup_method(method)
        self.config = self.get_base_config()
        self.es_context = ConnectionContext(self.config)

        self.index_client = self.es_context.indices_client()

        with self.es_context() as conn:
            self.connection = conn

        self.es_context.create_index(self.es_context.get_index_template())

    def teardown_method(self, method):
        # Clear the test indices.
        self.index_client.delete(self.es_context.get_index_template())
        super().teardown_method(method)

    def health_check(self):
        self.connection.cluster.health(wait_for_status="yellow",
                                       request_timeout=5)

    def get_url(self):
        """Returns the first url in the elasticsearch_urls list"""
        return self.config.elasticsearch_urls[0]

    def get_tuned_config(self, sources, extra_values=None):
        values_source = DEFAULT_VALUES.copy()
        if extra_values:
            values_source.update(extra_values)

        return super().get_tuned_config(sources, values_source)

    def get_base_config(self, cls=ConnectionContext, es_index=None):
        extra_values = None
        if es_index:
            extra_values = {
                "resource.elasticsearch.elasticsearch_index": es_index
            }

        return self.get_tuned_config(cls, extra_values=extra_values)

    def index_crash(self, processed_crash=None, raw_crash=None, crash_id=None):
        if crash_id is None:
            crash_id = str(uuid.UUID(int=random.getrandbits(128)))

        raw_crash = raw_crash or {}
        processed_crash = processed_crash or {}

        doc = {
            "crash_id": crash_id,
            "processed_crash": processed_crash,
            "raw_crash": raw_crash,
        }
        index_name = self.es_context.get_index_template()
        res = self.connection.index(
            index=index_name,
            doc_type=self.es_context.get_doctype(),
            id=crash_id,
            body=doc,
        )
        return res["_id"]

    def index_many_crashes(self,
                           number,
                           processed_crash=None,
                           raw_crash=None,
                           loop_field=None):
        processed_crash = processed_crash or {}
        raw_crash = raw_crash or {}

        actions = []
        for i in range(number):
            crash_id = str(uuid.UUID(int=random.getrandbits(128)))

            if loop_field is not None:
                processed_copy = processed_crash.copy()
                processed_copy[loop_field] = processed_crash[loop_field] % i
            else:
                processed_copy = processed_crash

            doc = {
                "crash_id": crash_id,
                "processed_crash": processed_copy,
                "raw_crash": raw_crash,
            }
            action = {
                "_index": self.es_context.get_index_template(),
                "_type": self.es_context.get_doctype(),
                "_id": crash_id,
                "_source": doc,
            }
            actions.append(action)

        bulk(client=self.connection, actions=actions)
        self.es_context.refresh()
Ejemplo n.º 2
0
class ElasticsearchTestCase(TestCaseWithConfig):
    """Base class for Elastic Search related unit tests"""
    def setup_method(self):
        super().setup_method()
        self.config = self.get_base_config()
        self.es_context = ConnectionContext(self.config)
        self.crashstorage = ESCrashStorage(
            config=self.get_tuned_config(ESCrashStorage))

        self.index_client = self.es_context.indices_client()
        self.conn = self.es_context.connection()

        # Delete everything there first
        for index_name in self.es_context.get_indices():
            print(f"setup: delete test index: {index_name}")
            self.es_context.delete_index(index_name)

        to_create = [
            self.es_context.get_index_for_date(utc_now()),
            self.es_context.get_index_for_date(utc_now() - timedelta(days=7)),
        ]
        for index_name in to_create:
            print(f"setup: creating index: {index_name}")
            self.es_context.create_index(index_name)

    def teardown_method(self):
        for index_name in self.es_context.get_indices():
            print(f"teardown: delete test index: {index_name}")
            self.es_context.delete_index(index_name)
        super().teardown_method()

    def health_check(self):
        self.conn.cluster.health(wait_for_status="yellow", request_timeout=5)

    def get_url(self):
        """Returns the first url in the elasticsearch_urls list"""
        return self.config.elasticsearch_urls[0]

    def get_tuned_config(self, sources, extra_values=None):
        values_source = DEFAULT_VALUES.copy()
        if extra_values:
            values_source.update(extra_values)

        return super().get_tuned_config(sources, values_source)

    def get_base_config(self, cls=ConnectionContext, es_index=None):
        extra_values = None
        if es_index:
            extra_values = {
                "resource.elasticsearch.elasticsearch_index": es_index
            }

        return self.get_tuned_config(cls, extra_values=extra_values)

    def index_crash(self,
                    processed_crash=None,
                    raw_crash=None,
                    crash_id=None,
                    refresh=True):
        """Index a single crash and refresh"""
        if crash_id is None:
            crash_id = str(uuid.UUID(int=random.getrandbits(128)))

        raw_crash = raw_crash or {}
        processed_crash = processed_crash or {}
        raw_crash["uuid"] = crash_id
        processed_crash["crash_id"] = crash_id
        processed_crash["uuid"] = crash_id

        self.crashstorage.save_processed_crash(raw_crash, processed_crash)

        if refresh:
            self.es_context.refresh()

        return crash_id

    def index_many_crashes(self,
                           number,
                           processed_crash=None,
                           raw_crash=None,
                           loop_field=None):
        """Index multiple crashes and refresh at the end"""
        processed_crash = processed_crash or {}
        raw_crash = raw_crash or {}

        crash_ids = []
        for i in range(number):
            if loop_field is not None:
                processed_copy = processed_crash.copy()
                processed_copy[loop_field] = processed_crash[loop_field] % i
            else:
                processed_copy = processed_crash

            crash_ids.append(
                self.index_crash(raw_crash=raw_crash,
                                 processed_crash=processed_copy,
                                 refresh=False))

        self.es_context.refresh()
        return crash_ids