f*ing UTs

Signed-off-by: Luka Peschke <[email protected]>
ToucanToco · Jan 9, 2025 · 08e0362 · 08e0362
1 parent a540411
commit 08e0362
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 9 deletions.
diff --git a/server/tests/backends/sql_translator_unit_tests/test_base_translator.py b/server/tests/backends/sql_translator_unit_tests/test_base_translator.py
@@ -3,13 +3,13 @@
 from zoneinfo import ZoneInfo
 
 import pytest
-from pypika import AliasedQuery, Case, Field, Order, Query, Schema, Table, analytics, functions
+from pypika import AliasedQuery, Case, Field, Query, Schema, Table, analytics, functions
 from pypika.enums import JoinType
 from pypika.queries import QueryBuilder
 from pypika.terms import LiteralValue, Term, ValueWrapper
 from pytest_mock import MockFixture
 
-from weaverbird.backends.pypika_translator.translators.base import DataTypeMapping, SQLTranslator
+from weaverbird.backends.pypika_translator.translators.base import DataTypeMapping, Order, SQLTranslator
 from weaverbird.backends.pypika_translator.translators.exceptions import (
     ForbiddenSQLStep,
     UnknownTableColumns,
@@ -199,7 +199,12 @@ def test_aggregate(base_translator: BaseTranslator, agg_type: str, default_step_
     agg_func = base_translator._get_aggregate_function(agg_type)
     field = Field(agg_field)
     agged = agg_func(field) if agg_func is not functions.Count else functions.Count("*")
-    expected_query = Query.from_(previous_step).groupby(field).orderby(agg_field).select(field, agged.as_(new_column))
+    expected_query = (
+        Query.from_(previous_step)
+        .groupby(field)
+        .orderby(agg_field, order=Order.ASC_NULLS_LAST)
+        .select(field, agged.as_(new_column))
+    )
 
     assert ctx.selectable.get_sql() == expected_query.get_sql()
 
@@ -231,7 +236,7 @@ def test_aggregate_with_original_granularity(
         .left_join(agg_query)
         .on_field(agg_field)
         .select(*original_select)
-        .orderby(agg_field)
+        .orderby(agg_field, order=Order.ASC_NULLS_LAST)
     )
 
     assert ctx.selectable.get_sql() == expected_query.get_sql()
@@ -544,7 +549,7 @@ def test_sort(base_translator: BaseTranslator, default_step_kwargs: dict[str, An
     step = steps.SortStep(columns=columns)
     ctx = base_translator.sort(step=step, columns=selected_columns, **default_step_kwargs)
 
-    expected_query = Query.from_(previous_step).select(*selected_columns).orderby(Field("name"), order=Order.asc)
+    expected_query = Query.from_(previous_step).select(*selected_columns).orderby(Field("name"), order=Order.ASC)
 
     assert ctx.selectable.get_sql() == expected_query.get_sql()
 
@@ -623,7 +628,12 @@ def test_uniquegroups(base_translator: BaseTranslator, default_step_kwargs: dict
     step = steps.UniqueGroupsStep(on=columns)
     ctx = base_translator.uniquegroups(step=step, columns=selected_columns, **default_step_kwargs)
 
-    expected_query = Query.from_(previous_step).select(Field(column)).groupby(Field(column)).orderby(Field(column))
+    expected_query = (
+        Query.from_(previous_step)
+        .select(Field(column))
+        .groupby(Field(column))
+        .orderby(Field(column), order=Order.ASC_NULLS_LAST)
+    )
 
     assert ctx.selectable.get_sql() == expected_query.get_sql()
 

diff --git a/server/tests/backends/sql_translator_unit_tests/test_base_translator_strings.py b/server/tests/backends/sql_translator_unit_tests/test_base_translator_strings.py
@@ -18,7 +18,7 @@
                 aggregations=[{"aggfunction": "count", "new_columns": ["beer_count"], "columns": ["name"]}],
             ),
         ],
-        'WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") SELECT "beer_kind",COUNT(*) "beer_count" FROM "__step_0_dummy__" GROUP BY "beer_kind" ORDER BY "beer_kind"',
+        'WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") SELECT "beer_kind",COUNT(*) "beer_count" FROM "__step_0_dummy__" GROUP BY "beer_kind" ORDER BY "beer_kind" ASC NULLS LAST',
     ),
     (
         [
@@ -29,7 +29,7 @@
                 keep_original_granularity=True,
             ),
         ],
-        'WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") SELECT "__step_0_dummy__"."price_per_l","__step_0_dummy__"."alcohol_degree","__step_0_dummy__"."name","__step_0_dummy__"."cost","__step_0_dummy__"."beer_kind","__step_0_dummy__"."volume_ml","__step_0_dummy__"."brewing_date","__step_0_dummy__"."nullable_name","sq0"."beer_count" FROM "__step_0_dummy__" LEFT JOIN (WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") SELECT "beer_kind",COUNT(*) "beer_count" FROM "__step_0_dummy__" GROUP BY "beer_kind") "sq0" ON "__step_0_dummy__"."beer_kind"="sq0"."beer_kind" ORDER BY "__step_0_dummy__"."beer_kind"',
+        'WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") SELECT "__step_0_dummy__"."price_per_l","__step_0_dummy__"."alcohol_degree","__step_0_dummy__"."name","__step_0_dummy__"."cost","__step_0_dummy__"."beer_kind","__step_0_dummy__"."volume_ml","__step_0_dummy__"."brewing_date","__step_0_dummy__"."nullable_name","sq0"."beer_count" FROM "__step_0_dummy__" LEFT JOIN (WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") SELECT "beer_kind",COUNT(*) "beer_count" FROM "__step_0_dummy__" GROUP BY "beer_kind") "sq0" ON "__step_0_dummy__"."beer_kind"="sq0"."beer_kind" ORDER BY "__step_0_dummy__"."beer_kind" ASC NULLS LAST',
     ),
     (
         [
@@ -47,7 +47,7 @@
             ),
             steps.AbsoluteValueStep(column="avg_price_per_l", new_column="avg_price_per_l_abs"),
         ],
-        'WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") ,__step_1_dummy__ AS (SELECT "beer_kind",COUNT(*) "beer_count",AVG("price_per_l") "avg_price_per_l" FROM "__step_0_dummy__" GROUP BY "beer_kind" ORDER BY "beer_kind") SELECT "beer_kind","beer_count","avg_price_per_l",ABS("avg_price_per_l") "avg_price_per_l_abs" FROM "__step_1_dummy__"',
+        'WITH __step_0_dummy__ AS (SELECT "price_per_l","alcohol_degree","name","cost","beer_kind","volume_ml","brewing_date","nullable_name" FROM "beers_tiny") ,__step_1_dummy__ AS (SELECT "beer_kind",COUNT(*) "beer_count",AVG("price_per_l") "avg_price_per_l" FROM "__step_0_dummy__" GROUP BY "beer_kind" ORDER BY "beer_kind" ASC NULLS LAST) SELECT "beer_kind","beer_count","avg_price_per_l",ABS("avg_price_per_l") "avg_price_per_l_abs" FROM "__step_1_dummy__"',
     ),
     (
         [