from sqlglot.dialects.dialect import Dialects
from tests.dialects.test_dialect import Validator


class TestDruid(Validator):
    dialect = "druid"

    def test_druid(self):
        self.validate_identity("SELECT CEIL(__time TO WEEK) FROM t")
        self.validate_identity("SELECT CEIL(col) FROM t")
        self.validate_identity("SELECT CEIL(price, 2) AS rounded_price FROM t")
        self.validate_identity("SELECT FLOOR(__time TO WEEK) FROM t")
        self.validate_identity("SELECT FLOOR(col) FROM t")
        self.validate_identity("SELECT FLOOR(price, 2) AS rounded_price FROM t")

        # validate across all dialects
        write = {dialect.value: "FLOOR(__time TO WEEK)" for dialect in Dialects}
        self.validate_all(
            "FLOOR(__time TO WEEK)",
            write=write,
        )