redash/tests/models/test_data_sources.py

import mock
from mock import patch
from tests import BaseTestCase

from redash.models import DataSource, Query, QueryResult
from redash.utils.configuration import ConfigurationContainer


class DataSourceTest(BaseTestCase):
    def test_get_schema(self):
        return_value = [{"name": "table", "columns": []}]

        with mock.patch(
            "redash.query_runner.pg.PostgreSQL.get_schema"
        ) as patched_get_schema:
            patched_get_schema.return_value = return_value

            schema = self.factory.data_source.get_schema()

            self.assertEqual(return_value, schema)

    def test_get_schema_uses_cache(self):
        return_value = [{"name": "table", "columns": []}]
        with mock.patch(
            "redash.query_runner.pg.PostgreSQL.get_schema"
        ) as patched_get_schema:
            patched_get_schema.return_value = return_value

            self.factory.data_source.get_schema()
            schema = self.factory.data_source.get_schema()

            self.assertEqual(return_value, schema)
            self.assertEqual(patched_get_schema.call_count, 1)

    def test_get_schema_skips_cache_with_refresh_true(self):
        return_value = [{"name": "table", "columns": []}]
        with mock.patch(
            "redash.query_runner.pg.PostgreSQL.get_schema"
        ) as patched_get_schema:
            patched_get_schema.return_value = return_value

            self.factory.data_source.get_schema()
            new_return_value = [{"name": "new_table", "columns": []}]
            patched_get_schema.return_value = new_return_value
            schema = self.factory.data_source.get_schema(refresh=True)

            self.assertEqual(new_return_value, schema)
            self.assertEqual(patched_get_schema.call_count, 2)

    def test_schema_sorter(self):
        input_data = [
            {"name": "zoo", "columns": ["is_zebra", "is_snake", "is_cow"]},
            {
                "name": "all_terain_vehicle",
                "columns": ["has_wheels", "has_engine", "has_all_wheel_drive"],
            },
        ]

        expected_output = [
            {
                "name": "all_terain_vehicle",
                "columns": ["has_all_wheel_drive", "has_engine", "has_wheels"],
            },
            {"name": "zoo", "columns": ["is_cow", "is_snake", "is_zebra"]},
        ]

        real_output = self.factory.data_source._sort_schema(input_data)

        self.assertEqual(real_output, expected_output)

    def test_model_uses_schema_sorter(self):
        orig_schema = [
            {"name": "zoo", "columns": ["is_zebra", "is_snake", "is_cow"]},
            {
                "name": "all_terain_vehicle",
                "columns": ["has_wheels", "has_engine", "has_all_wheel_drive"],
            },
        ]

        sorted_schema = [
            {
                "name": "all_terain_vehicle",
                "columns": ["has_all_wheel_drive", "has_engine", "has_wheels"],
            },
            {"name": "zoo", "columns": ["is_cow", "is_snake", "is_zebra"]},
        ]

        with mock.patch(
            "redash.query_runner.pg.PostgreSQL.get_schema"
        ) as patched_get_schema:
            patched_get_schema.return_value = orig_schema

            out_schema = self.factory.data_source.get_schema()

            self.assertEqual(out_schema, sorted_schema)


class TestDataSourceCreate(BaseTestCase):
    def test_adds_data_source_to_default_group(self):
        data_source = DataSource.create_with_group(
            org=self.factory.org,
            name="test",
            options=ConfigurationContainer.from_json('{"dbname": "test"}'),
            type="pg",
        )
        self.assertIn(self.factory.org.default_group.id, data_source.groups)


class TestDataSourceIsPaused(BaseTestCase):
    def test_returns_false_by_default(self):
        self.assertFalse(self.factory.data_source.paused)

    def test_persists_selection(self):
        self.factory.data_source.pause()
        self.assertTrue(self.factory.data_source.paused)

        self.factory.data_source.resume()
        self.assertFalse(self.factory.data_source.paused)

    def test_allows_setting_reason(self):
        reason = "Some good reason."
        self.factory.data_source.pause(reason)
        self.assertTrue(self.factory.data_source.paused)
        self.assertEqual(self.factory.data_source.pause_reason, reason)

    def test_resume_clears_reason(self):
        self.factory.data_source.pause("Reason")
        self.factory.data_source.resume()
        self.assertEqual(self.factory.data_source.pause_reason, None)

    def test_reason_is_none_by_default(self):
        self.assertEqual(self.factory.data_source.pause_reason, None)


class TestDataSourceDelete(BaseTestCase):
    def test_deletes_the_data_source(self):
        data_source = self.factory.create_data_source()
        data_source.delete()

        self.assertIsNone(DataSource.query.get(data_source.id))

    def test_sets_queries_data_source_to_null(self):
        data_source = self.factory.create_data_source()
        query = self.factory.create_query(data_source=data_source)

        data_source.delete()
        self.assertIsNone(DataSource.query.get(data_source.id))
        self.assertIsNone(Query.query.get(query.id).data_source_id)

    def test_deletes_child_models(self):
        data_source = self.factory.create_data_source()
        self.factory.create_query_result(data_source=data_source)
        self.factory.create_query(
            data_source=data_source,
            latest_query_data=self.factory.create_query_result(data_source=data_source),
        )

        data_source.delete()
        self.assertIsNone(DataSource.query.get(data_source.id))
        self.assertEqual(
            0, QueryResult.query.filter(QueryResult.data_source == data_source).count()
        )

    @patch("redash.redis_connection.delete")
    def test_deletes_schema(self, mock_redis):
        data_source = self.factory.create_data_source()
        data_source.delete()

        mock_redis.assert_called_with(data_source._schema_key)