move sdk join tests

getsentry · Jun 27, 2023 · ab3e534 · ab3e534
1 parent 556ba66
commit ab3e534
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 80 deletions.
diff --git a/tests/test_join_api.py b/tests/test_join_api.py
@@ -6,6 +6,17 @@
 
 import pytest
 import simplejson as json
+from snuba_sdk import (
+ Column,
+ Condition,
+ Entity,
+ Function,
+ Join,
+ Op,
+ Query,
+ Relationship,
+ Request,
+)
 
 from snuba.datasets.entities.entity_key import EntityKey
 from snuba.datasets.entities.factory import get_entity
@@ -222,3 +233,79 @@ def test_invalid_columns_composite_query(
  # assert json.loads(response.data)["error"]["message"] == error_message
 
  assert response.status_code == 500
+
+ def test_sdk_join_query(self) -> None:
+ ev = Entity("events", "ev")
+ gm = Entity("groupedmessage", "gm")
+ join = Join([Relationship(ev, "grouped", gm)])
+ query = (
+ Query(join)
+ .set_select(
+ [
+ Column("group_id", ev),
+ Column("status", gm),
+ Function("avg", [Column("retention_days", ev)], "avg"),
+ ]
+ )
+ .set_groupby([Column("group_id", ev), Column("status", gm)])
+ .set_where(
+ [
+ Condition(Column("project_id", ev), Op.EQ, self.project_id),
+ Condition(Column("project_id", gm), Op.EQ, self.project_id),
+ Condition(Column("timestamp", ev), Op.GTE, self.base_time),
+ Condition(Column("timestamp", ev), Op.LT, self.next_time),
+ ]
+ )
+ )
+
+ request = Request(
+ dataset="discover",
+ query=query,
+ app_id="default",
+ tenant_ids={"referrer": "r", "organization_id": 123},
+ )
+ response = self.post("/discover/snql", data=json.dumps(request.to_dict()))
+ data = json.loads(response.data)
+
+ assert response.status_code == 200
+ assert data["data"] == []
+
+ def test_array_condition_unpack_in_join_query(self) -> None:
+ ev = Entity("events", "ev")
+ gm = Entity("groupedmessage", "gm")
+ join = Join([Relationship(ev, "grouped", gm)])
+ query = (
+ Query(join)
+ .set_select(
+ [
+ Column("group_id", ev),
+ Column("status", gm),
+ Function("avg", [Column("retention_days", ev)], "avg"),
+ ]
+ )
+ .set_groupby([Column("group_id", ev), Column("status", gm)])
+ .set_where(
+ [
+ Condition(Column("project_id", ev), Op.EQ, self.project_id),
+ Condition(Column("project_id", gm), Op.EQ, self.project_id),
+ Condition(Column("timestamp", ev), Op.GTE, self.base_time),
+ Condition(Column("timestamp", ev), Op.LT, self.next_time),
+ Condition(
+ Column("exception_stacks.type", ev), Op.LIKE, "Arithmetic%"
+ ),
+ ]
+ )
+ )
+
+ request = Request(
+ dataset="discover",
+ query=query,
+ app_id="default",
+ tenant_ids={"referrer": "r", "organization_id": 123},
+ )
+ request.flags.debug = True
+ response = self.post("/discover/snql", data=json.dumps(request.to_dict()))
+ data = json.loads(response.data)
+
+ assert response.status_code == 200
+ assert data["data"] == []
diff --git a/tests/test_snql_sdk_api.py b/tests/test_snql_sdk_api.py
@@ -11,12 +11,10 @@
  Entity,
  Function,
  Identifier,
- Join,
  Lambda,
  Op,
  OrderBy,
  Query,
- Relationship,
  Request,
 )
 
@@ -29,8 +27,6 @@
 from tests.fixtures import get_raw_event, get_raw_transaction
 from tests.helpers import write_unprocessed_events
 
-pytest.skip(allow_module_level=True, reason="Dataset no longer exists")
-
 
 @pytest.mark.clickhouse_db
 @pytest.mark.redis_db
@@ -133,42 +129,6 @@ def test_sessions_query(self) -> None:
  assert response.status_code == 200
  assert data["data"] == []
 
- def test_join_query(self) -> None:
- ev = Entity("events", "ev")
- gm = Entity("groupedmessage", "gm")
- join = Join([Relationship(ev, "grouped", gm)])
- query = (
- Query(join)
- .set_select(
- [
- Column("group_id", ev),
- Column("status", gm),
- Function("avg", [Column("retention_days", ev)], "avg"),
- ]
- )
- .set_groupby([Column("group_id", ev), Column("status", gm)])
- .set_where(
- [
- Condition(Column("project_id", ev), Op.EQ, self.project_id),
- Condition(Column("project_id", gm), Op.EQ, self.project_id),
- Condition(Column("timestamp", ev), Op.GTE, self.base_time),
- Condition(Column("timestamp", ev), Op.LT, self.next_time),
- ]
- )
- )
-
- request = Request(
- dataset="discover",
- query=query,
- app_id="default",
- tenant_ids={"referrer": "r", "organization_id": 123},
- )
- response = self.post("/discover/snql", data=json.dumps(request.to_dict()))
- data = json.loads(response.data)
-
- assert response.status_code == 200
- assert data["data"] == []
-
  def test_sub_query(self) -> None:
  inner_query = (
  Query(Entity("discover_events"))
@@ -291,46 +251,6 @@ def test_tags_in_groupby(self) -> None:
  data = json.loads(response.data)
  assert response.status_code == 200, data
 
- def test_array_condition_unpack_in_join_query(self) -> None:
- ev = Entity("events", "ev")
- gm = Entity("groupedmessage", "gm")
- join = Join([Relationship(ev, "grouped", gm)])
- query = (
- Query(join)
- .set_select(
- [
- Column("group_id", ev),
- Column("status", gm),
- Function("avg", [Column("retention_days", ev)], "avg"),
- ]
- )
- .set_groupby([Column("group_id", ev), Column("status", gm)])
- .set_where(
- [
- Condition(Column("project_id", ev), Op.EQ, self.project_id),
- Condition(Column("project_id", gm), Op.EQ, self.project_id),
- Condition(Column("timestamp", ev), Op.GTE, self.base_time),
- Condition(Column("timestamp", ev), Op.LT, self.next_time),
- Condition(
- Column("exception_stacks.type", ev), Op.LIKE, "Arithmetic%"
- ),
- ]
- )
- )
-
- request = Request(
- dataset="discover",
- query=query,
- app_id="default",
- tenant_ids={"referrer": "r", "organization_id": 123},
- )
- request.flags.debug = True
- response = self.post("/discover/snql", data=json.dumps(request.to_dict()))
- data = json.loads(response.data)
-
- assert response.status_code == 200
- assert data["data"] == []
-
  def test_escape_edge_cases(self) -> None:
  query = (
  Query(Entity("events"))