redash/tests/query_runner/test_mongodb.py

import datetime
from unittest import TestCase

from pytz import utc
from freezegun import freeze_time

from redash.query_runner.mongodb import (
    parse_query_json,
    parse_results,
    _get_column_by_name,
)
from redash.utils import json_dumps, parse_human_time


class TestParseQueryJson(TestCase):
    def test_ignores_non_isodate_fields(self):
        query = {"test": 1, "test_list": ["a", "b", "c"], "test_dict": {"a": 1, "b": 2}}

        query_data = parse_query_json(json_dumps(query))
        self.assertDictEqual(query_data, query)

    def test_parses_isodate_fields(self):
        query = {
            "test": 1,
            "test_list": ["a", "b", "c"],
            "test_dict": {"a": 1, "b": 2},
            "testIsoDate": 'ISODate("2014-10-03T00:00")',
        }

        query_data = parse_query_json(json_dumps(query))

        self.assertEqual(
            query_data["testIsoDate"], datetime.datetime(2014, 10, 3, 0, 0)
        )

    def test_parses_isodate_in_nested_fields(self):
        query = {
            "test": 1,
            "test_list": ["a", "b", "c"],
            "test_dict": {"a": 1, "b": {"date": 'ISODate("2014-10-04T00:00")'}},
            "testIsoDate": 'ISODate("2014-10-03T00:00")',
        }

        query_data = parse_query_json(json_dumps(query))

        self.assertEqual(
            query_data["testIsoDate"], datetime.datetime(2014, 10, 3, 0, 0)
        )
        self.assertEqual(
            query_data["test_dict"]["b"]["date"], datetime.datetime(2014, 10, 4, 0, 0)
        )

    def test_handles_nested_fields(self):
        # https://github.com/getredash/redash/issues/597
        query = {
            "collection": "bus",
            "aggregate": [
                {
                    "$geoNear": {
                        "near": {
                            "type": "Point",
                            "coordinates": [-22.910079, -43.205161],
                        },
                        "maxDistance": 100000000,
                        "distanceField": "dist.calculated",
                        "includeLocs": "dist.location",
                        "spherical": True,
                    }
                }
            ],
        }

        query_data = parse_query_json(json_dumps(query))

        self.assertDictEqual(query, query_data)

    def test_supports_extended_json_types(self):
        query = {
            "test": 1,
            "test_list": ["a", "b", "c"],
            "test_dict": {"a": 1, "b": 2},
            "testIsoDate": 'ISODate("2014-10-03T00:00")',
            "test$date": {"$date": "2014-10-03T00:00:00.0"},
            "test$undefined": {"$undefined": None},
        }
        query_data = parse_query_json(json_dumps(query))
        self.assertEqual(query_data["test$undefined"], None)
        self.assertEqual(
            query_data["test$date"],
            datetime.datetime(2014, 10, 3, 0, 0).replace(tzinfo=utc),
        )

    @freeze_time("2019-01-01 12:00:00")
    def test_supports_relative_timestamps(self):
        query = {"ts": {"$humanTime": "1 hour ago"}}

        one_hour_ago = parse_human_time("1 hour ago")
        query_data = parse_query_json(json_dumps(query))
        self.assertEqual(query_data["ts"], one_hour_ago)


class TestMongoResults(TestCase):
    def test_parses_regular_results(self):
        raw_results = [
            {"column": 1, "column2": "test"},
            {"column": 2, "column2": "test", "column3": "hello"},
        ]
        rows, columns = parse_results(raw_results)

        for i, row in enumerate(rows):
            self.assertDictEqual(row, raw_results[i])

        self.assertIsNotNone(_get_column_by_name(columns, "column"))
        self.assertIsNotNone(_get_column_by_name(columns, "column2"))
        self.assertIsNotNone(_get_column_by_name(columns, "column3"))

    def test_parses_nested_results(self):
        raw_results = [
            {"column": 1, "column2": "test", "nested": {"a": 1, "b": "str"}},
            {
                "column": 2,
                "column2": "test",
                "column3": "hello",
                "nested": {"a": 2, "b": "str2", "c": "c"},
            },
        ]

        rows, columns = parse_results(raw_results)

        self.assertDictEqual(
            rows[0], {"column": 1, "column2": "test", "nested.a": 1, "nested.b": "str"}
        )
        self.assertDictEqual(
            rows[1],
            {
                "column": 2,
                "column2": "test",
                "column3": "hello",
                "nested.a": 2,
                "nested.b": "str2",
                "nested.c": "c",
            },
        )

        self.assertIsNotNone(_get_column_by_name(columns, "column"))
        self.assertIsNotNone(_get_column_by_name(columns, "column2"))
        self.assertIsNotNone(_get_column_by_name(columns, "column3"))
        self.assertIsNotNone(_get_column_by_name(columns, "nested.a"))
        self.assertIsNotNone(_get_column_by_name(columns, "nested.b"))
        self.assertIsNotNone(_get_column_by_name(columns, "nested.c"))