From 4c8009103187913f4ac52266eab74b0545114e75 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Tue, 14 Jan 2025 16:30:26 -0500 Subject: [PATCH 1/9] feat: improve ta_storage.bigquery - generalize the way bigquery accepts query parameters - change type of params arg in bigquery_service to sequence - feat: add upload_id field to ta_testrun protobuf - add flags_hash field to ta_testrun protobuf - create new testid generation function - add test_id to ta_testrun proto - add flaky_failure to testrun protobuf - handle flaky failures in ta_storage.bq - create sql queries for reading from bq - write tests for ta_storage.bq aggregate queries --- database/models/reports.py | 2 +- generated_proto/testrun/ta_testrun_pb2.py | 27 +- protobuf/ta_testrun.proto | 5 + requirements.in | 1 + requirements.txt | 6 +- services/bigquery.py | 18 +- services/tests/test_bigquery.py | 28 ++ ta_storage/base.py | 1 + ta_storage/bq.py | 264 ++++++++++++- ta_storage/pg.py | 2 +- ta_storage/tests/test_bq.py | 374 +++++++++++++++++- ta_storage/utils.py | 28 ++ ...cessorTask__ta_processor_task_call__2.json | 20 +- 13 files changed, 736 insertions(+), 40 deletions(-) create mode 100644 ta_storage/utils.py diff --git a/database/models/reports.py b/database/models/reports.py index 922fceed3..d727c494a 100644 --- a/database/models/reports.py +++ b/database/models/reports.py @@ -117,7 +117,7 @@ class Upload(CodecovBaseModel, MixinBaseClass): upload_type_id = Column(types.Integer) @cached_property - def flag_names(self): + def flag_names(self) -> list[str]: return [f.flag_name for f in self.flags] diff --git a/generated_proto/testrun/ta_testrun_pb2.py b/generated_proto/testrun/ta_testrun_pb2.py index ba133248f..b9da9d04e 100644 --- a/generated_proto/testrun/ta_testrun_pb2.py +++ b/generated_proto/testrun/ta_testrun_pb2.py @@ -4,32 +4,35 @@ # source: ta_testrun.proto # Protobuf Python Version: 5.29.2 """Generated protocol buffer code.""" - from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import runtime_version as _runtime_version from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder - _runtime_version.ValidateProtobufRuntimeVersion( - _runtime_version.Domain.PUBLIC, 5, 29, 2, "", "ta_testrun.proto" + _runtime_version.Domain.PUBLIC, + 5, + 29, + 2, + '', + 'ta_testrun.proto' ) # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n\x10ta_testrun.proto"\xda\x02\n\x07TestRun\x12\x11\n\ttimestamp\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tclassname\x18\x03 \x01(\t\x12\x11\n\ttestsuite\x18\x04 \x01(\t\x12\x15\n\rcomputed_name\x18\x05 \x01(\t\x12!\n\x07outcome\x18\x06 \x01(\x0e\x32\x10.TestRun.Outcome\x12\x17\n\x0f\x66\x61ilure_message\x18\x07 \x01(\t\x12\x18\n\x10\x64uration_seconds\x18\x08 \x01(\x02\x12\x0e\n\x06repoid\x18\n \x01(\x03\x12\x12\n\ncommit_sha\x18\x0b \x01(\t\x12\x13\n\x0b\x62ranch_name\x18\x0c \x01(\t\x12\r\n\x05\x66lags\x18\r \x03(\t\x12\x10\n\x08\x66ilename\x18\x0e \x01(\t\x12\x11\n\tframework\x18\x0f \x01(\t".\n\x07Outcome\x12\n\n\x06PASSED\x10\x00\x12\n\n\x06\x46\x41ILED\x10\x01\x12\x0b\n\x07SKIPPED\x10\x02' -) + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x10ta_testrun.proto\"\xa4\x03\n\x07TestRun\x12\x11\n\ttimestamp\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tclassname\x18\x03 \x01(\t\x12\x11\n\ttestsuite\x18\x04 \x01(\t\x12\x15\n\rcomputed_name\x18\x05 \x01(\t\x12!\n\x07outcome\x18\x06 \x01(\x0e\x32\x10.TestRun.Outcome\x12\x17\n\x0f\x66\x61ilure_message\x18\x07 \x01(\t\x12\x18\n\x10\x64uration_seconds\x18\x08 \x01(\x02\x12\x0e\n\x06repoid\x18\n \x01(\x03\x12\x12\n\ncommit_sha\x18\x0b \x01(\t\x12\x13\n\x0b\x62ranch_name\x18\x0c \x01(\t\x12\r\n\x05\x66lags\x18\r \x03(\t\x12\x10\n\x08\x66ilename\x18\x0e \x01(\t\x12\x11\n\tframework\x18\x0f \x01(\t\x12\x11\n\tupload_id\x18\x10 \x01(\x03\x12\x12\n\nflags_hash\x18\x11 \x01(\x0c\x12\x0f\n\x07test_id\x18\x12 \x01(\x0c\"@\n\x07Outcome\x12\n\n\x06PASSED\x10\x00\x12\n\n\x06\x46\x41ILED\x10\x01\x12\x0b\n\x07SKIPPED\x10\x02\x12\x10\n\x0c\x46LAKY_FAILED\x10\x03') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "ta_testrun_pb2", _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'ta_testrun_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: - DESCRIPTOR._loaded_options = None - _globals["_TESTRUN"]._serialized_start = 21 - _globals["_TESTRUN"]._serialized_end = 367 - _globals["_TESTRUN_OUTCOME"]._serialized_start = 321 - _globals["_TESTRUN_OUTCOME"]._serialized_end = 367 + DESCRIPTOR._loaded_options = None + _globals['_TESTRUN']._serialized_start=21 + _globals['_TESTRUN']._serialized_end=441 + _globals['_TESTRUN_OUTCOME']._serialized_start=377 + _globals['_TESTRUN_OUTCOME']._serialized_end=441 # @@protoc_insertion_point(module_scope) diff --git a/protobuf/ta_testrun.proto b/protobuf/ta_testrun.proto index c0530b875..7b460c43a 100644 --- a/protobuf/ta_testrun.proto +++ b/protobuf/ta_testrun.proto @@ -11,6 +11,7 @@ message TestRun { PASSED = 0; FAILED = 1; SKIPPED = 2; + FLAKY_FAILED = 3; } optional Outcome outcome = 6; @@ -27,4 +28,8 @@ message TestRun { optional string filename = 14; optional string framework = 15; + + optional int64 upload_id = 16; + optional bytes flags_hash = 17; + optional bytes test_id = 18; } diff --git a/requirements.in b/requirements.in index 8533a5625..fbc593be1 100644 --- a/requirements.in +++ b/requirements.in @@ -20,6 +20,7 @@ grpcio>=1.66.2 httpx jinja2>=3.1.3 lxml>=5.3.0 +mmh3>=5.0.1 mock multidict>=6.1.0 openai diff --git a/requirements.txt b/requirements.txt index 4265f9cd5..995223df8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -208,8 +208,10 @@ markupsafe==2.1.3 # via jinja2 minio==7.1.13 # via shared -mmh3==4.0.1 - # via shared +mmh3==5.0.1 + # via + # -r requirements.in + # shared mock==4.0.3 # via -r requirements.in monotonic==1.5 diff --git a/services/bigquery.py b/services/bigquery.py index 4da0de60a..1fda004f8 100644 --- a/services/bigquery.py +++ b/services/bigquery.py @@ -1,5 +1,5 @@ from types import ModuleType -from typing import Dict, List, cast +from typing import Dict, List, Sequence, cast import polars as pl from google.api_core import retry @@ -108,7 +108,17 @@ def write( self.write_client.batch_commit_write_streams(batch_commit_write_streams_request) - def query(self, query: str, params: dict | None = None) -> List[Dict]: + def query( + self, + query: str, + params: Sequence[ + bigquery.ScalarQueryParameter + | bigquery.RangeQueryParameter + | bigquery.ArrayQueryParameter + | bigquery.StructQueryParameter + ] + | None = None, + ) -> List[Dict]: """Execute a BigQuery SQL query and return results. Try not to write INSERT statements and use the write method instead. @@ -125,9 +135,7 @@ def query(self, query: str, params: dict | None = None) -> List[Dict]: job_config = bigquery.QueryJobConfig() if params: - job_config.query_parameters = [ - bigquery.ScalarQueryParameter(k, "STRING", v) for k, v in params.items() - ] + job_config.query_parameters = params row_iterator = self.client.query_and_wait( query, job_config=job_config, retry=retry.Retry(deadline=30) diff --git a/services/tests/test_bigquery.py b/services/tests/test_bigquery.py index 9a9198d69..f2c861d19 100644 --- a/services/tests/test_bigquery.py +++ b/services/tests/test_bigquery.py @@ -3,6 +3,7 @@ import polars as pl import pytest +from google.cloud.bigquery import ScalarQueryParameter import generated_proto.testrun.ta_testrun_pb2 as ta_testrun_pb2 from services.bigquery import BigQueryService @@ -63,6 +64,33 @@ def test_bigquery_service(): assert {row["id"] for row in results} == {1, 2} +sql = """ +WITH sample_data AS ( + SELECT * FROM UNNEST([ + STRUCT(TIMESTAMP '2025-01-01T00:00:00Z' AS timestamp, 1 AS id, 'name' AS name), + STRUCT(TIMESTAMP '2024-12-30T00:00:00Z' AS timestamp, 2 AS id, 'name2' AS name) + ]) +) +SELECT * FROM sample_data where id = @id +""" + + +@pytest.mark.skip(reason="This test requires being run using actual working creds") +def test_bigquery_service_params(): + bigquery_service = BigQueryService(gcp_config) + + results = bigquery_service.query( + sql, params=[ScalarQueryParameter("id", "INT64", 2)] + ) + + assert len(results) == 1 + assert {row["timestamp"] for row in results} == { + datetime.fromisoformat("2024-12-30T00:00:00Z"), + } + assert {row["name"] for row in results} == {"name2"} + assert {row["id"] for row in results} == {2} + + @pytest.mark.skip(reason="This test requires being run using actual working creds") def test_bigquery_service_polars(): bigquery_service = BigQueryService(gcp_config) diff --git a/ta_storage/base.py b/ta_storage/base.py index ddeb5558b..690e3fefa 100644 --- a/ta_storage/base.py +++ b/ta_storage/base.py @@ -18,5 +18,6 @@ def write_testruns( upload: Upload, framework: str | None, testruns: list[test_results_parser.Testrun], + flaky_test_set: set[str], ): pass diff --git a/ta_storage/bq.py b/ta_storage/bq.py index 22828c584..7d9658c75 100644 --- a/ta_storage/bq.py +++ b/ta_storage/bq.py @@ -4,23 +4,138 @@ from typing import Literal, cast import test_results_parser +from google.cloud.bigquery import ArrayQueryParameter, ScalarQueryParameter from shared.config import get_config import generated_proto.testrun.ta_testrun_pb2 as ta_testrun_pb2 from database.models.reports import Upload from services.bigquery import get_bigquery_service from ta_storage.base import TADriver +from ta_storage.utils import calc_flags_hash, calc_test_id -DATASET_NAME: str = cast( - str, get_config("services", "bigquery", "dataset_name", default="codecov_prod") +RANKED_DATA = """ +ranked_data AS ( + SELECT + *, + ROW_NUMBER() OVER ( + PARTITION BY + name, + classname, + testsuite, + flags_hash + ORDER BY timestamp DESC + ) AS row_num + FROM + `{PROJECT_ID}.{DATASET_NAME}.{TESTRUN_TABLE_NAME}` + WHERE + repoid = @repoid + AND commit_sha = @commit_sha ) +""" -TESTRUN_TABLE_NAME: str = cast( - str, get_config("services", "bigquery", "testrun_table_name", default="testruns") +LATEST_INSTANCES = """ +latest_instances AS ( + SELECT + * + FROM + ranked_data + WHERE + row_num = 1 ) +""" +PR_COMMENT_AGG = """ +SELECT + * +FROM ( + SELECT + commit_sha, + outcome + FROM + latest_instances +) PIVOT ( + COUNT(*) AS ct + FOR outcome IN ( + 0 as passed, + 1 as failed, + 2 as skipped, + 3 as flaky_failed + ) +) +""" + +PR_COMMENT_FAIL = """ +SELECT + computed_name, + failure_message, + flags +FROM + latest_instances +WHERE + outcome = 1 +""" + +TESTRUNS_FOR_UPLOAD = """ +SELECT + DATE_BUCKET(timestamp, INTERVAL 1 DAY) AS date, + test_id, + outcome, + branch_name, +FROM + `{PROJECT_ID}.{DATASET_NAME}.{TESTRUN_TABLE_NAME}` +WHERE + upload_id = @upload_id + AND ( + outcome = 1 + OR outcome = 3 + OR test_id IN UNNEST(@test_ids) + ) +""" + +ANALYTICS_BASE = """ +analytics_base AS ( + SELECT * + FROM `{PROJECT_ID}.{DATASET_NAME}.{TESTRUN_TABLE_NAME}` + WHERE repoid = @repoid + AND timestamp BETWEEN + (CURRENT_DATE - INTERVAL @interval_start) AND + (CURRENT_DATE - INTERVAL @interval_end) +) +""" + +ANALYTICS_BRANCH = """ +analytics_base AS ( + SELECT * + FROM `{PROJECT_ID}.{DATASET_NAME}.{TESTRUN_TABLE_NAME}` + WHERE repoid = @repoid + AND branch_name = @branch + AND timestamp BETWEEN + (CURRENT_TIMESTAMP() - INTERVAL @interval_start DAY) AND + (CURRENT_TIMESTAMP() - INTERVAL @interval_end DAY) +) +""" -def outcome_to_int( +ANALYTICS = """ +SELECT + name, + classname, + testsuite, + ANY_VALUE(computed_name) AS computed_name, + COUNT(DISTINCT IF(outcome = 1 OR outcome = 3, commit_sha, NULL)) AS cwf, + AVG(duration_seconds) AS avg_duration, + MAX_BY(duration_seconds, timestamp) AS last_duration, + SUM(IF(outcome = 0, 1, 0)) AS pass_count, + SUM(IF(outcome = 1, 1, 0)) AS fail_count, + SUM(IF(outcome = 2, 1, 0)) AS skip_count, + SUM(IF(outcome = 3, 1, 0)) AS flaky_fail_count, + MAX(timestamp) AS updated_at, + ARRAY_AGG(DISTINCT unnested_flags) AS flags +FROM analytics_base, UNNEST(flags) AS unnested_flags +GROUP BY name, classname, testsuite +""" + + +def outcome_to_enum( outcome: Literal["pass", "skip", "failure", "error"], ) -> ta_testrun_pb2.TestRun.Outcome: match outcome: @@ -35,6 +150,26 @@ def outcome_to_int( class BQDriver(TADriver): + def __init__(self, flaky_test_set: set[bytes] | None = None): + self.flaky_test_set = flaky_test_set or {} + self.bq_service = get_bigquery_service() + + self.project_id: str = cast( + str, get_config("services", "gcp", "project_id", default="codecov-prod") + ) + + self.dataset_name: str = cast( + str, + get_config("services", "bigquery", "dataset_name", default="codecov_prod"), + ) + + self.testrun_table_name: str = cast( + str, + get_config( + "services", "bigquery", "testrun_table_name", default="testruns" + ), + ) + def write_testruns( self, timestamp: int | None, @@ -45,15 +180,21 @@ def write_testruns( framework: str | None, testruns: list[test_results_parser.Testrun], ): - bq_service = get_bigquery_service() - if timestamp is None: timestamp = int(datetime.now().timestamp() * 1000000) - flag_names = upload.flag_names + flag_names: list[str] = upload.flag_names testruns_pb: list[bytes] = [] + flags_hash = calc_flags_hash(flag_names) + for t in testruns: + test_id = calc_test_id(t["name"], t["classname"], t["testsuite"]) + if test_id in self.flaky_test_set and t["outcome"] == "failure": + outcome = ta_testrun_pb2.TestRun.Outcome.FLAKY_FAILED + else: + outcome = outcome_to_enum(t["outcome"]) + test_run = ta_testrun_pb2.TestRun( timestamp=timestamp, repoid=repo_id, @@ -65,12 +206,113 @@ def write_testruns( name=t["name"], testsuite=t["testsuite"], computed_name=t["computed_name"], - outcome=outcome_to_int(t["outcome"]), + outcome=outcome, failure_message=t["failure_message"], duration_seconds=t["duration"], filename=t["filename"], + upload_id=upload.id_, + flags_hash=flags_hash, + test_id=test_id, ) testruns_pb.append(test_run.SerializeToString()) - flag_names = upload.flag_names - bq_service.write(DATASET_NAME, TESTRUN_TABLE_NAME, ta_testrun_pb2, testruns_pb) + self.bq_service.write( + self.dataset_name, self.testrun_table_name, ta_testrun_pb2, testruns_pb + ) + + def pr_comment_agg( + self, + repoid: int, + commit_sha: str, + ): + query = f""" + WITH + {RANKED_DATA.format( + PROJECT_ID=self.project_id, + DATASET_NAME=self.dataset_name, + TESTRUN_TABLE_NAME=self.testrun_table_name, + )}, + {LATEST_INSTANCES} + {PR_COMMENT_AGG} + """ + return self.bq_service.query( + query, + [ + ScalarQueryParameter("repoid", "INT64", repoid), + ScalarQueryParameter("commit_sha", "STRING", commit_sha), + ], + ) + + def pr_comment_fail(self, repoid: int, commit_sha: str): + query = f""" + WITH + {RANKED_DATA.format( + PROJECT_ID=self.project_id, + DATASET_NAME=self.dataset_name, + TESTRUN_TABLE_NAME=self.testrun_table_name, + )}, + {LATEST_INSTANCES} + {PR_COMMENT_FAIL} + """ + return self.bq_service.query( + query, + [ + ScalarQueryParameter("repoid", "INT64", repoid), + ScalarQueryParameter("commit_sha", "STRING", commit_sha), + ], + ) + + def testruns_for_upload(self, upload_id: int, test_ids: list[bytes]): + query = f""" + {TESTRUNS_FOR_UPLOAD.format( + PROJECT_ID=self.project_id, + DATASET_NAME=self.dataset_name, + TESTRUN_TABLE_NAME=self.testrun_table_name, + )} + """ + return self.bq_service.query( + query, + [ + ScalarQueryParameter("upload_id", "INT64", upload_id), + ArrayQueryParameter("test_ids", "BYTES", test_ids), + ], + ) + + def analytics( + self, + repoid: int, + interval_start: int = 30, # for convention we want the start to be the larger number of days + interval_end: int = 0, + branch: str | None = None, + ): + if branch: + query = f""" + WITH + {ANALYTICS_BRANCH.format( + PROJECT_ID=self.project_id, + DATASET_NAME=self.dataset_name, + TESTRUN_TABLE_NAME=self.testrun_table_name, + )} + {ANALYTICS} + """ + else: + query = f""" + WITH + {ANALYTICS_BASE.format( + PROJECT_ID=self.project_id, + DATASET_NAME=self.dataset_name, + TESTRUN_TABLE_NAME=self.testrun_table_name, + )} + {ANALYTICS} + """ + + params = [ + ScalarQueryParameter("repoid", "INT64", repoid), + ScalarQueryParameter("interval_start", "INT64", interval_start), + ScalarQueryParameter("interval_end", "INT64", interval_end), + ] + + if branch: + params.append(ScalarQueryParameter("branch", "STRING", branch)) + + return self.bq_service.query(query, params) diff --git a/ta_storage/pg.py b/ta_storage/pg.py index c829a3adf..762efd8bd 100644 --- a/ta_storage/pg.py +++ b/ta_storage/pg.py @@ -283,7 +283,7 @@ def save_test_instances(db_session: Session, test_instance_data: list[dict]): class PGDriver(TADriver): - def __init__(self, db_session: Session, flaky_test_set: set): + def __init__(self, db_session: Session, flaky_test_set: set[str]): self.db_session = db_session self.flaky_test_set = flaky_test_set diff --git a/ta_storage/tests/test_bq.py b/ta_storage/tests/test_bq.py index d78840c25..208190098 100644 --- a/ta_storage/tests/test_bq.py +++ b/ta_storage/tests/test_bq.py @@ -1,14 +1,16 @@ from __future__ import annotations -from datetime import datetime +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch import pytest import test_results_parser +import time_machine import generated_proto.testrun.ta_testrun_pb2 as ta_testrun_pb2 from database.tests.factories import RepositoryFlagFactory, UploadFactory -from ta_storage.bq import DATASET_NAME, TESTRUN_TABLE_NAME, BQDriver +from ta_storage.bq import BQDriver +from ta_storage.utils import calc_flags_hash, calc_test_id @pytest.fixture @@ -77,10 +79,12 @@ def test_bigquery_driver(dbsession, mock_bigquery_service): test_data, ) + flags_hash = calc_flags_hash(upload.flag_names) + # Verify the BigQuery service was called correctly mock_bigquery_service.write.assert_called_once_with( - DATASET_NAME, - TESTRUN_TABLE_NAME, + bq.dataset_name, + bq.testrun_table_name, ta_testrun_pb2, [ ta_testrun_pb2.TestRun( @@ -98,6 +102,9 @@ def test_bigquery_driver(dbsession, mock_bigquery_service): framework="pytest", branch_name=upload.report.commit.branch, flags=["flag1", "flag2"], + upload_id=upload.id_, + flags_hash=flags_hash, + test_id=calc_test_id("test_name", "test_class", "test_suite"), ).SerializeToString(), ta_testrun_pb2.TestRun( timestamp=timestamp, @@ -114,6 +121,365 @@ def test_bigquery_driver(dbsession, mock_bigquery_service): framework="pytest", branch_name=upload.report.commit.branch, flags=["flag1", "flag2"], + upload_id=upload.id_, + flags_hash=flags_hash, + test_id=calc_test_id("test_name2", "test_class2", "test_suite2"), ).SerializeToString(), ], ) + + +def populate_pr_comment_testruns(bq: BQDriver): + testruns = [] + + for i in range(3): + upload = UploadFactory() + upload.report.commit.commitid = "abcde" + upload.report.commit.branch = "feature_branch" + upload.report.commit.repoid = 2 + upload.flags.append(RepositoryFlagFactory(flag_name=f"flag_{i}")) + + for j in range(3): + name = f"test_{j}" + classname = f"class_{j}" + testsuite = "suite_feature" + + testrun: test_results_parser.Testrun = { + "name": name, + "classname": classname, + "testsuite": testsuite, + "duration": float(j % 5), + "outcome": "pass" if j % 2 == 0 else "failure", + "filename": None, + "computed_name": f"pr_computed_name_{j}", + "failure_message": None if j % 2 == 0 else "hi", + "build_url": None, + } + + testruns.append(testrun) + + bq.write_testruns( + None, 2, "abcde", "feature_branch", upload, "pytest", testruns + ) + + +@pytest.mark.skip(reason="need creds") +def test_bq_pr_comment(): + bq = BQDriver() + + if ( + bq.bq_service.query( + "select * from `test_dataset.testruns` where repoid = 2 limit 1" + ) + == [] + ): + populate_pr_comment_testruns(bq) + + pr_agg = bq.pr_comment_agg(repoid=2, commit_sha="abcde") + assert pr_agg == [ + { + "commit_sha": "abcde", + "ct_passed": 6, + "ct_failed": 3, + "ct_skipped": 0, + "ct_flaky_failed": 0, + } + ] + + pr_fail = bq.pr_comment_fail(repoid=2, commit_sha="abcde") + assert len(pr_fail) == 3 + assert {t["computed_name"] for t in pr_fail} == { + "pr_computed_name_1", + } + assert {t["failure_message"] for t in pr_fail} == {"hi"} + assert {tuple(t["flags"]) for t in pr_fail} == { + ("flag_1",), + ("flag_2",), + ("flag_0",), + } + + +def populate_testruns_for_upload_testruns(dbsession, bq: BQDriver): + testruns = [] + + upload = UploadFactory() + upload.id_ = 1 + dbsession.add(upload) + dbsession.flush() + + testruns: list[test_results_parser.Testrun] = [ + { # this test is flaky failure + "name": "test_0", + "classname": "class_0", + "testsuite": "suite_upload", + "duration": 0.0, + "outcome": "failure", + "filename": None, + "computed_name": "upload_computed_name_0", + "failure_message": None, + "build_url": None, + }, + { # this test is just a failure + "name": "test_1", + "classname": "class_1", + "testsuite": "suite_upload", + "duration": 0.0, + "outcome": "failure", + "filename": None, + "computed_name": "upload_computed_name_1", + "failure_message": None, + "build_url": None, + }, + { # this test is a pass but also flaky + "name": "test_2", + "classname": "class_2", + "testsuite": "suite_upload", + "duration": 0.0, + "outcome": "pass", + "filename": None, + "computed_name": "upload_computed_name_2", + "failure_message": None, + "build_url": None, + }, + { # this test should be ignored + "name": "test_3", + "classname": "class_3", + "testsuite": "suite_upload", + "duration": 0.0, + "outcome": "pass", + "filename": None, + "computed_name": "upload_computed_name_3", + "failure_message": None, + "build_url": None, + }, + ] + + bq.write_testruns(None, 3, "abcde", "feature_branch", upload, "pytest", testruns) + + +@pytest.mark.skip(reason="need creds") +def test_bq_testruns_for_upload(dbsession): + bq = BQDriver( + { + calc_test_id("test_0", "class_0", "suite_upload"), + calc_test_id("test_2", "class_2", "suite_upload"), + } + ) + + if ( + bq.bq_service.query( + "select * from `test_dataset.testruns` where repoid = 3 limit 1" + ) + == [] + ): + populate_testruns_for_upload_testruns(dbsession, bq) + + testruns_for_upload = bq.testruns_for_upload( + upload_id=1, + test_ids=[ + calc_test_id("test_0", "class_0", "suite_upload"), + calc_test_id("test_2", "class_2", "suite_upload"), + ], + ) + + assert {t["test_id"] for t in testruns_for_upload} == { + calc_test_id("test_0", "class_0", "suite_upload"), + calc_test_id("test_2", "class_2", "suite_upload"), + calc_test_id("test_1", "class_1", "suite_upload"), + } + + assert {t["outcome"] for t in testruns_for_upload} == {3, 1, 0} + + +def populate_analytics_testruns(bq: BQDriver): + upload_0 = UploadFactory() + upload_0.report.commit.commitid = "abcde" + upload_0.report.commit.branch = "feature_branch" + upload_0.report.commit.repoid = 1 + upload_0.flags.append(RepositoryFlagFactory(flag_name="flag_0")) + + upload_1 = UploadFactory() + upload_1.report.commit.commitid = "abcde" + upload_1.report.commit.branch = "feature_branch" + upload_1.report.commit.repoid = 1 + upload_1.flags.append(RepositoryFlagFactory(flag_name="flag_1")) + + testruns: list[test_results_parser.Testrun] = [ + { + "name": "interval_start", + "classname": "class_0", + "testsuite": "suite_upload", + "duration": 20000.0, + "outcome": "failure", + "filename": None, + "computed_name": "upload_computed_name_0", + "failure_message": None, + "build_url": None, + }, + ] + + timestamp = int((datetime.now() - timedelta(days=50)).timestamp() * 1000000) + + bq.write_testruns( + timestamp, 1, "interval_start", "feature_branch", upload_0, "pytest", testruns + ) + + testruns: list[test_results_parser.Testrun] = [ + { + "name": "interval_end", + "classname": "class_0", + "testsuite": "suite_upload", + "duration": 20000.0, + "outcome": "failure", + "filename": None, + "computed_name": "upload_computed_name_0", + "failure_message": None, + "build_url": None, + }, + ] + + timestamp = int((datetime.now() - timedelta(days=1)).timestamp() * 1000000) + + bq.write_testruns( + timestamp, 1, "interval_end", "feature_branch", upload_0, "pytest", testruns + ) + + testruns: list[test_results_parser.Testrun] = [ + { + "name": "test_0", + "classname": "class_0", + "testsuite": "suite_upload", + "duration": 10.0, + "outcome": "failure", + "filename": None, + "computed_name": "upload_computed_name_0", + "failure_message": None, + "build_url": None, + }, + { + "name": "test_1", + "classname": "class_1", + "testsuite": "suite_upload", + "duration": 10.0, + "outcome": "pass", + "filename": None, + "computed_name": "upload_computed_name_1", + "failure_message": None, + "build_url": None, + }, + ] + + timestamp = int((datetime.now() - timedelta(days=20)).timestamp() * 1000000) + + bq.write_testruns( + timestamp, 1, "commit_1", "feature_branch", upload_0, "pytest", testruns + ) + + testruns: list[test_results_parser.Testrun] = [ + { + "name": "test_1", + "classname": "class_1", + "testsuite": "suite_upload", + "duration": 10.0, + "outcome": "failure", + "filename": None, + "computed_name": "upload_computed_name_1", + "failure_message": None, + "build_url": None, + }, + ] + + timestamp = int((datetime.now() - timedelta(days=20)).timestamp() * 1000000) + + bq.write_testruns( + timestamp, 1, "commit_1", "feature_branch", upload_1, "pytest", testruns + ) + + bq = BQDriver( + { + calc_test_id("test_1", "class_1", "suite_upload"), + } + ) + + testruns: list[test_results_parser.Testrun] = [ + { + "name": "test_0", + "classname": "class_0", + "testsuite": "suite_upload", + "duration": 20.0, + "outcome": "pass", + "filename": None, + "computed_name": "upload_computed_name_0", + "failure_message": None, + "build_url": None, + }, + { + "name": "test_1", + "classname": "class_1", + "testsuite": "suite_upload", + "duration": 10.0, + "outcome": "failure", + "filename": None, + "computed_name": "upload_computed_name_1", + "failure_message": None, + "build_url": None, + }, + ] + + timestamp = int((datetime.now() - timedelta(days=10)).timestamp() * 1000000) + + bq.write_testruns( + timestamp, 1, "commit_2", "feature_branch", upload_1, "pytest", testruns + ) + + +@pytest.mark.skip(reason="need creds") +@time_machine.travel(datetime.now(tz=timezone.utc), tick=False) +def test_bq_analytics(): + bq = BQDriver() + + if ( + bq.bq_service.query( + "select * from `test_dataset.testruns` where repoid = 1 limit 1" + ) + == [] + ): + populate_analytics_testruns(bq) + + testruns_for_upload = bq.analytics(1, 30, 7, "feature_branch") + + assert sorted( + [(x | {"flags": sorted(x["flags"])}) for x in testruns_for_upload], + key=lambda x: x["name"], + ) == [ + { + "name": "test_0", + "classname": "class_0", + "testsuite": "suite_upload", + "computed_name": "upload_computed_name_0", + "cwf": 1, + "avg_duration": 15.0, + "last_duration": 20.0, + "pass_count": 1, + "fail_count": 1, + "skip_count": 0, + "flaky_fail_count": 0, + "updated_at": datetime.now(tz=timezone.utc) - timedelta(days=10), + "flags": ["flag_0", "flag_1"], + }, + { + "name": "test_1", + "classname": "class_1", + "testsuite": "suite_upload", + "computed_name": "upload_computed_name_1", + "cwf": 2, + "avg_duration": 10.0, + "last_duration": 10.0, + "pass_count": 1, + "fail_count": 1, + "skip_count": 0, + "flaky_fail_count": 1, + "updated_at": datetime.now(tz=timezone.utc) - timedelta(days=10), + "flags": ["flag_0", "flag_1"], + }, + ] diff --git a/ta_storage/utils.py b/ta_storage/utils.py new file mode 100644 index 000000000..83ad866ca --- /dev/null +++ b/ta_storage/utils.py @@ -0,0 +1,28 @@ +from sys import byteorder + +import mmh3 +import sentry_sdk + + +def calc_test_id(name: str, classname: str, testsuite: str) -> bytes: + h = mmh3.mmh3_x64_128() # assumes we're running on x64 machines + h.update(testsuite.encode("utf-8")) + h.update(classname.encode("utf-8")) + h.update(name.encode("utf-8")) + test_id_hash = h.digest() + + return test_id_hash + + +def calc_flags_hash(flags: list[str]) -> bytes | None: + flags_str = " ".join(flags) # we know that flags cannot contain spaces + + # returns a tuple of two int64 values + # we only need the first one + flags_hash, _ = mmh3.hash64(flags_str) + try: + flags_hash_bytes = flags_hash.to_bytes(8, byteorder) + return flags_hash_bytes + except OverflowError as e: # this should never happen because hash64 should always return 2 64 bit ints + sentry_sdk.capture_exception(e) + return None diff --git a/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json b/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json index 73abd5112..936abe615 100644 --- a/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json +++ b/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json @@ -9,7 +9,10 @@ "duration_seconds": 0.001, "repoid": "1", "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", - "framework": "Pytest" + "framework": "Pytest", + "upload_id": "1", + "flags_hash": "AAAAAAAAAAA=", + "test_id": "S/K2VdzrrehI4hnoZNsPVg==" }, { "timestamp": "1735689600000000", @@ -22,7 +25,10 @@ "duration_seconds": 0.001, "repoid": "1", "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", - "framework": "Pytest" + "framework": "Pytest", + "upload_id": "1", + "flags_hash": "AAAAAAAAAAA=", + "test_id": "CVU2jNUNOkOrl6/lJdK0nw==" }, { "timestamp": "1735689600000000", @@ -34,7 +40,10 @@ "duration_seconds": 0.0, "repoid": "1", "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", - "framework": "Pytest" + "framework": "Pytest", + "upload_id": "1", + "flags_hash": "AAAAAAAAAAA=", + "test_id": "UDBibp0NWEToP72TpCn1xg==" }, { "timestamp": "1735689600000000", @@ -46,6 +55,9 @@ "duration_seconds": 0.001, "repoid": "1", "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", - "framework": "Pytest" + "framework": "Pytest", + "upload_id": "1", + "flags_hash": "AAAAAAAAAAA=", + "test_id": "VE2yD2IYxdSbTvGB6XCJPA==" } ] From 45f646381e94444bf9b2c21fdf9853c4b6162727 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 24 Jan 2025 16:10:12 -0500 Subject: [PATCH 2/9] fix: address feedback --- ta_storage/utils.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/ta_storage/utils.py b/ta_storage/utils.py index 83ad866ca..b83c57f7a 100644 --- a/ta_storage/utils.py +++ b/ta_storage/utils.py @@ -1,5 +1,3 @@ -from sys import byteorder - import mmh3 import sentry_sdk @@ -15,13 +13,13 @@ def calc_test_id(name: str, classname: str, testsuite: str) -> bytes: def calc_flags_hash(flags: list[str]) -> bytes | None: - flags_str = " ".join(flags) # we know that flags cannot contain spaces + flags_str = " ".join(sorted(flags)) # we know that flags cannot contain spaces # returns a tuple of two int64 values # we only need the first one - flags_hash, _ = mmh3.hash64(flags_str) + flags_hash, _ = mmh3.hash64(flags_str, signed=False) try: - flags_hash_bytes = flags_hash.to_bytes(8, byteorder) + flags_hash_bytes = flags_hash.to_bytes(8) return flags_hash_bytes except OverflowError as e: # this should never happen because hash64 should always return 2 64 bit ints sentry_sdk.capture_exception(e) From 95f393afb06cd13f6547dc8d3ecba24924925101 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 24 Jan 2025 16:50:35 -0500 Subject: [PATCH 3/9] fix lint --- generated_proto/testrun/ta_testrun_pb2.py | 27 ++++++++++------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/generated_proto/testrun/ta_testrun_pb2.py b/generated_proto/testrun/ta_testrun_pb2.py index b9da9d04e..04e4f6247 100644 --- a/generated_proto/testrun/ta_testrun_pb2.py +++ b/generated_proto/testrun/ta_testrun_pb2.py @@ -4,35 +4,32 @@ # source: ta_testrun.proto # Protobuf Python Version: 5.29.2 """Generated protocol buffer code.""" + from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import runtime_version as _runtime_version from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder + _runtime_version.ValidateProtobufRuntimeVersion( - _runtime_version.Domain.PUBLIC, - 5, - 29, - 2, - '', - 'ta_testrun.proto' + _runtime_version.Domain.PUBLIC, 5, 29, 2, "", "ta_testrun.proto" ) # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() - - -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x10ta_testrun.proto\"\xa4\x03\n\x07TestRun\x12\x11\n\ttimestamp\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tclassname\x18\x03 \x01(\t\x12\x11\n\ttestsuite\x18\x04 \x01(\t\x12\x15\n\rcomputed_name\x18\x05 \x01(\t\x12!\n\x07outcome\x18\x06 \x01(\x0e\x32\x10.TestRun.Outcome\x12\x17\n\x0f\x66\x61ilure_message\x18\x07 \x01(\t\x12\x18\n\x10\x64uration_seconds\x18\x08 \x01(\x02\x12\x0e\n\x06repoid\x18\n \x01(\x03\x12\x12\n\ncommit_sha\x18\x0b \x01(\t\x12\x13\n\x0b\x62ranch_name\x18\x0c \x01(\t\x12\r\n\x05\x66lags\x18\r \x03(\t\x12\x10\n\x08\x66ilename\x18\x0e \x01(\t\x12\x11\n\tframework\x18\x0f \x01(\t\x12\x11\n\tupload_id\x18\x10 \x01(\x03\x12\x12\n\nflags_hash\x18\x11 \x01(\x0c\x12\x0f\n\x07test_id\x18\x12 \x01(\x0c\"@\n\x07Outcome\x12\n\n\x06PASSED\x10\x00\x12\n\n\x06\x46\x41ILED\x10\x01\x12\x0b\n\x07SKIPPED\x10\x02\x12\x10\n\x0c\x46LAKY_FAILED\x10\x03') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x10ta_testrun.proto"\xa4\x03\n\x07TestRun\x12\x11\n\ttimestamp\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tclassname\x18\x03 \x01(\t\x12\x11\n\ttestsuite\x18\x04 \x01(\t\x12\x15\n\rcomputed_name\x18\x05 \x01(\t\x12!\n\x07outcome\x18\x06 \x01(\x0e\x32\x10.TestRun.Outcome\x12\x17\n\x0f\x66\x61ilure_message\x18\x07 \x01(\t\x12\x18\n\x10\x64uration_seconds\x18\x08 \x01(\x02\x12\x0e\n\x06repoid\x18\n \x01(\x03\x12\x12\n\ncommit_sha\x18\x0b \x01(\t\x12\x13\n\x0b\x62ranch_name\x18\x0c \x01(\t\x12\r\n\x05\x66lags\x18\r \x03(\t\x12\x10\n\x08\x66ilename\x18\x0e \x01(\t\x12\x11\n\tframework\x18\x0f \x01(\t\x12\x11\n\tupload_id\x18\x10 \x01(\x03\x12\x12\n\nflags_hash\x18\x11 \x01(\x0c\x12\x0f\n\x07test_id\x18\x12 \x01(\x0c"@\n\x07Outcome\x12\n\n\x06PASSED\x10\x00\x12\n\n\x06\x46\x41ILED\x10\x01\x12\x0b\n\x07SKIPPED\x10\x02\x12\x10\n\x0c\x46LAKY_FAILED\x10\x03' +) _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'ta_testrun_pb2', _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "ta_testrun_pb2", _globals) if not _descriptor._USE_C_DESCRIPTORS: - DESCRIPTOR._loaded_options = None - _globals['_TESTRUN']._serialized_start=21 - _globals['_TESTRUN']._serialized_end=441 - _globals['_TESTRUN_OUTCOME']._serialized_start=377 - _globals['_TESTRUN_OUTCOME']._serialized_end=441 + DESCRIPTOR._loaded_options = None + _globals["_TESTRUN"]._serialized_start = 21 + _globals["_TESTRUN"]._serialized_end = 441 + _globals["_TESTRUN_OUTCOME"]._serialized_start = 377 + _globals["_TESTRUN_OUTCOME"]._serialized_end = 441 # @@protoc_insertion_point(module_scope) From 30b426f3398c3dbe7dcdf8ce659de5bf36c253b7 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 24 Jan 2025 16:53:53 -0500 Subject: [PATCH 4/9] address feedback i missed --- ta_storage/utils.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/ta_storage/utils.py b/ta_storage/utils.py index b83c57f7a..0c3dd3e42 100644 --- a/ta_storage/utils.py +++ b/ta_storage/utils.py @@ -1,5 +1,4 @@ import mmh3 -import sentry_sdk def calc_test_id(name: str, classname: str, testsuite: str) -> bytes: @@ -18,9 +17,5 @@ def calc_flags_hash(flags: list[str]) -> bytes | None: # returns a tuple of two int64 values # we only need the first one flags_hash, _ = mmh3.hash64(flags_str, signed=False) - try: - flags_hash_bytes = flags_hash.to_bytes(8) - return flags_hash_bytes - except OverflowError as e: # this should never happen because hash64 should always return 2 64 bit ints - sentry_sdk.capture_exception(e) - return None + flags_hash_bytes = flags_hash.to_bytes(8) + return flags_hash_bytes From 6c6004bf3b8a1d7bae78f2fa5253b15a0ca0a5a9 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 17 Jan 2025 14:39:07 -0500 Subject: [PATCH 5/9] improve BQDriver's handling of flags hash Improve TADriver interface - add some more methods to the TADriver - implement a base constructor - modify the write_testruns interface - implement all methods in BQ and PG - improve BQ and PG tests - modify use of TADriver interface in processor and finishers - update django settings to include new settings - TODO: modify requirements to suitable shared version - create ta_utils to replace test_results in the future - the reason for this is that we want a slightly different implementation of the test results notifier for the new TA pipeline --- database/tests/factories/reports.py | 16 + django_scaffold/settings.py | 13 + services/processing/flake_processing.py | 18 +- services/ta_utils.py | 373 ++++++++++ services/tests/test_ta_utils.py | 317 ++++++++ ta_storage/base.py | 59 +- ta_storage/bq.py | 340 ++++++++- ta_storage/pg.py | 484 +++++++++++- .../tests/snapshots/bq__analytics__0.txt | 31 + .../bq__analytics_with_branch__0.txt | 32 + .../snapshots/bq__cache_analytics__0.json | 40 + .../snapshots/bq__cache_analytics__1.txt | 32 + .../tests/snapshots/bq__pr_comment_agg__0.txt | 51 ++ .../snapshots/bq__pr_comment_agg__1.json | 7 + .../snapshots/bq__pr_comment_fail__0.txt | 45 ++ .../snapshots/bq__pr_comment_fail__1.json | 12 + .../tests/snapshots/bq__write_flakes__0.txt | 24 + .../tests/snapshots/bq__write_flakes__1.json | 12 + .../snapshots/bq__write_testruns__0.json | 22 + .../bq__write_testruns_with_flake__0.json | 23 + ta_storage/tests/test_bq.py | 702 ++++++++---------- ta_storage/tests/test_pg.py | 115 ++- ta_storage/tests/test_pg_django.py | 272 +++++++ ta_storage/utils.py | 6 + tasks/ta_processor.py | 56 +- ...cessorTask__ta_processor_task_call__2.json | 4 - tasks/tests/unit/test_ta_processor_task.py | 11 +- worker.sh | 4 + 28 files changed, 2627 insertions(+), 494 deletions(-) create mode 100644 services/ta_utils.py create mode 100644 services/tests/test_ta_utils.py create mode 100644 ta_storage/tests/snapshots/bq__analytics__0.txt create mode 100644 ta_storage/tests/snapshots/bq__analytics_with_branch__0.txt create mode 100644 ta_storage/tests/snapshots/bq__cache_analytics__0.json create mode 100644 ta_storage/tests/snapshots/bq__cache_analytics__1.txt create mode 100644 ta_storage/tests/snapshots/bq__pr_comment_agg__0.txt create mode 100644 ta_storage/tests/snapshots/bq__pr_comment_agg__1.json create mode 100644 ta_storage/tests/snapshots/bq__pr_comment_fail__0.txt create mode 100644 ta_storage/tests/snapshots/bq__pr_comment_fail__1.json create mode 100644 ta_storage/tests/snapshots/bq__write_flakes__0.txt create mode 100644 ta_storage/tests/snapshots/bq__write_flakes__1.json create mode 100644 ta_storage/tests/snapshots/bq__write_testruns__0.json create mode 100644 ta_storage/tests/snapshots/bq__write_testruns_with_flake__0.json create mode 100644 ta_storage/tests/test_pg_django.py diff --git a/database/tests/factories/reports.py b/database/tests/factories/reports.py index b5859dd0f..227b70f7e 100644 --- a/database/tests/factories/reports.py +++ b/database/tests/factories/reports.py @@ -7,12 +7,14 @@ Flake, RepositoryFlag, Test, + TestInstance, TestResultReportTotals, ) from database.tests.factories.core import ( CompareCommitFactory, ReportFactory, RepositoryFactory, + UploadFactory, ) @@ -43,6 +45,20 @@ class Meta: repository = factory.SubFactory(RepositoryFactory) +class TestInstanceFactory(factory.Factory): + class Meta: + model = TestInstance + + test = factory.SubFactory(TestFactory) + upload = factory.SubFactory(UploadFactory) + + outcome = "pass" + duration_seconds = 1.5 + commitid = factory.SelfAttribute("upload.report.commit.commitid") + branch = factory.SelfAttribute("upload.report.commit.branch") + repoid = factory.SelfAttribute("upload.report.commit.repository.repoid") + + class FlakeFactory(factory.Factory): class Meta: model = Flake diff --git a/django_scaffold/settings.py b/django_scaffold/settings.py index 7ff275eb4..9fa8208a4 100644 --- a/django_scaffold/settings.py +++ b/django_scaffold/settings.py @@ -12,11 +12,23 @@ if "timeseries" in DATABASES: DATABASES["timeseries"]["AUTOCOMMIT"] = False +if "test_analytics" in DATABASES: + DATABASES["test_analytics"]["AUTOCOMMIT"] = False + IS_DEV = os.getenv("RUN_ENV") == "DEV" IS_ENTERPRISE = os.getenv("RUN_ENV") == "ENTERPRISE" GCS_BUCKET_NAME = get_config("services", "minio", "bucket", default="codecov") +BIGQUERY_WRITE_ENABLED = ( + get_config("services", "bigquery", "write_enabled", default=False) + and "test_analytics" in DATABASES +) + +BIGQUERY_READ_ENABLED = ( + get_config("services", "bigquery", "read_enabled", default=False) + and "test_analytics" in DATABASES +) # Application definition INSTALLED_APPS = [ # dependencies @@ -40,6 +52,7 @@ "shared.django_apps.profiling", "shared.django_apps.reports", "shared.django_apps.staticanalysis", + "shared.django_apps.test_analytics", ] TELEMETRY_VANILLA_DB = "default" diff --git a/services/processing/flake_processing.py b/services/processing/flake_processing.py index 7d231b744..5d72fc34a 100644 --- a/services/processing/flake_processing.py +++ b/services/processing/flake_processing.py @@ -27,6 +27,17 @@ def process_flake_for_repo_commit( state__in=["processed", "v2_finished"], ).all() + process_flakes_for_uploads(repo_id, [upload for upload in uploads]) + + log.info( + "Successfully processed flakes", + extra=dict(repoid=repo_id, commit=commit_id), + ) + + return {"successful": True} + + +def process_flakes_for_uploads(repo_id: int, uploads: list[ReportSession]): curr_flakes = fetch_curr_flakes(repo_id) new_flakes: dict[str, Flake] = dict() @@ -92,13 +103,6 @@ def process_flake_for_repo_commit( upload.save() django_transaction.commit() - log.info( - "Successfully processed flakes", - extra=dict(repoid=repo_id, commit=commit_id), - ) - - return {"successful": True} - def get_test_instances( upload: ReportSession, diff --git a/services/ta_utils.py b/services/ta_utils.py new file mode 100644 index 000000000..f3425b651 --- /dev/null +++ b/services/ta_utils.py @@ -0,0 +1,373 @@ +import logging +from dataclasses import dataclass +from hashlib import sha256 + +from shared.plan.constants import FREE_PLAN_REPRESENTATIONS, TEAM_PLAN_REPRESENTATIONS +from shared.yaml import UserYaml +from sqlalchemy import desc, distinct, func +from sqlalchemy.orm import Session, joinedload + +from database.models import ( + Commit, + Repository, + TestInstance, + Upload, + UploadError, +) +from helpers.notifier import BaseNotifier +from rollouts import FLAKY_TEST_DETECTION +from services.license import requires_license +from services.urls import get_members_url, get_test_analytics_url +from services.yaml import read_yaml_field + +log = logging.getLogger(__name__) + + +def generate_flags_hash(flag_names: list[str]) -> str: + return sha256((" ".join(sorted(flag_names))).encode("utf-8")).hexdigest() + + +def generate_test_id(repoid, testsuite, name, flags_hash): + return sha256( + (" ".join([str(x) for x in [repoid, testsuite, name, flags_hash]])).encode( + "utf-8" + ) + ).hexdigest() + + +def latest_failures_for_commit( + db_session: Session, repo_id: int, commit_sha: str +) -> list[TestInstance]: + """ + This will result in a SQL query that looks something like this: + + SELECT DISTINCT ON (rti.test_id) rti.id, ... + FROM reports_testinstance rti + JOIN reports_upload ru ON ru.id = rti.upload_id + LEFT OUTER JOIN reports_test rt ON rt.id = rti.test_id + WHERE ... + ORDER BY rti.test_id, ru.created_at DESC + + The goal of this query is to return: + > the latest test instance failure for each unique test based on upload creation time + + The `DISTINCT ON` test_id with the order by test_id, enforces that we are only fetching one test instance for each test. + + The ordering by `upload.create_at DESC` enforces that we get the latest test instance for that unique test. + """ + + return ( + db_session.query(TestInstance) + .join(TestInstance.upload) + .options(joinedload(TestInstance.test)) + .filter(TestInstance.repoid == repo_id, TestInstance.commitid == commit_sha) + .filter(TestInstance.outcome.in_(["failure", "error"])) + .order_by(TestInstance.test_id) + .order_by(desc(Upload.created_at)) + .distinct(TestInstance.test_id) + .all() + ) + + +def get_test_summary_for_commit( + db_session: Session, repo_id: int, commit_sha: str +) -> dict[str, int]: + return dict( + db_session.query( + TestInstance.outcome, func.count(distinct(TestInstance.test_id)) + ) + .filter(TestInstance.repoid == repo_id, TestInstance.commitid == commit_sha) + .group_by(TestInstance.outcome) + .all() + ) + + +@dataclass +class FlakeInfo: + failed: int + count: int + + +@dataclass +class TestFailure: + display_name: str + failure_message: str | None + duration_seconds: float + build_url: str | None + flake_info: FlakeInfo | None = None + + +@dataclass +class TestResultsNotificationPayload: + failed: int + passed: int + skipped: int + regular_failures: list[TestFailure] | None = None + flaky_failures: list[TestFailure] | None = None + + +def wrap_in_details(summary: str, content: str): + result = f"
{summary}\n{content}\n
" + return result + + +def make_quoted(content: str) -> str: + lines = content.splitlines() + result = "\n".join("> " + line for line in lines) + return f"\n{result}\n" + + +def properly_backtick(content: str) -> str: + max_backtick_count = 0 + curr_backtick_count = 0 + prev_char = None + for char in content: + if char == "`": + curr_backtick_count += 1 + else: + curr_backtick_count = 0 + + if curr_backtick_count > max_backtick_count: + max_backtick_count = curr_backtick_count + + backticks = "`" * (max_backtick_count + 1) + return f"{backticks}python\n{content}\n{backticks}" + + +def wrap_in_code(content: str) -> str: + if "```" in content: + return properly_backtick(content) + else: + return f"\n```python\n{content}\n```\n" + + +def display_duration(f: float) -> str: + before_dot, after_dot = str(f).split(".") + if len(before_dot) > 3: + return before_dot + else: + return f"{f:.3g}" + + +def generate_failure_info( + fail: TestFailure, +): + if fail.failure_message is not None: + failure_message = fail.failure_message + else: + failure_message = "No failure message available" + + failure_message = wrap_in_code(failure_message) + if fail.build_url: + return f"{failure_message}\n[View]({fail.build_url}) the CI Build" + else: + return failure_message + + +def generate_view_test_analytics_line(commit: Commit) -> str: + repo = commit.repository + test_analytics_url = get_test_analytics_url(repo, commit) + return f"\nTo view more test analytics, go to the [Test Analytics Dashboard]({test_analytics_url})\n:loudspeaker: Thoughts on this report? [Let us know!](https://github.com/codecov/feedback/issues/304)" + + +def messagify_failure( + failure: TestFailure, +) -> str: + test_name = wrap_in_code(failure.display_name.replace("\x1f", " ")) + formatted_duration = display_duration(failure.duration_seconds) + stack_trace_summary = f"Stack Traces | {formatted_duration}s run time" + stack_trace = wrap_in_details( + stack_trace_summary, + make_quoted(generate_failure_info(failure)), + ) + return make_quoted(f"{test_name}\n{stack_trace}") + + +def messagify_flake( + failure: TestFailure, +) -> str: + test_name = wrap_in_code(failure.display_name.replace("\x1f", " ")) + formatted_duration = display_duration(failure.duration_seconds) + assert failure.flake_info is not None + flake_rate = failure.flake_info.failed / failure.flake_info.count * 100 + flake_rate_section = f"**Flake rate in main:** {flake_rate:.2f}% (Passed {failure.flake_info.count - failure.flake_info.failed} times, Failed {failure.flake_info.failed} times)" + stack_trace_summary = f"Stack Traces | {formatted_duration}s run time" + stack_trace = wrap_in_details( + stack_trace_summary, + make_quoted(generate_failure_info(failure)), + ) + return make_quoted(f"{test_name}\n{flake_rate_section}\n{stack_trace}") + + +def specific_error_message(upload_error: UploadError) -> str: + title = f"### :x: {upload_error.error_code}" + if upload_error.error_code == "Unsupported file format": + description = "\n".join( + [ + "Upload processing failed due to unsupported file format. Please review the parser error message:", + f"`{upload_error.error_params['error_message']}`", + "For more help, visit our [troubleshooting guide](https://docs.codecov.com/docs/test-analytics#troubleshooting).", + ] + ) + elif upload_error.error_code == "File not found": + description = "\n".join( + [ + "No result to display due to the CLI not being able to find the file.", + "Please ensure the file contains `junit` in the name and automated file search is enabled,", + "or the desired file specified by the `file` and `search_dir` arguments of the CLI.", + ] + ) + else: + raise ValueError("Unrecognized error code") + message = [ + title, + make_quoted(description), + ] + return "\n".join(message) + + +@dataclass +class TestResultsNotifier(BaseNotifier): + payload: TestResultsNotificationPayload | None = None + error: UploadError | None = None + + def build_message(self) -> str: + if self.payload is None: + raise ValueError("Payload passed to notifier is None, cannot build message") + + message = [] + + if self.error: + message.append(specific_error_message(self.error)) + + if self.error and ( + self.payload.regular_failures or self.payload.flaky_failures + ): + message.append("---") + + if self.payload.regular_failures or self.payload.flaky_failures: + message.append(f"### :x: {self.payload.failed} Tests Failed:") + + completed = self.payload.failed + self.payload.passed + + message += [ + "| Tests completed | Failed | Passed | Skipped |", + "|---|---|---|---|", + f"| {completed} | {self.payload.failed} | {self.payload.passed} | {self.payload.skipped} |", + ] + + if self.payload.regular_failures: + failures = sorted( + self.payload.regular_failures, + key=lambda x: (x.duration_seconds, x.display_name), + )[:3] + failure_content = [ + f"{messagify_failure(failure)}" for failure in failures + ] + + top_3_failed_section = wrap_in_details( + f"View the top {min(3, len(failures))} failed tests by shortest run time", + "\n".join(failure_content), + ) + + message.append(top_3_failed_section) + + if self.payload.regular_failures and self.payload.flaky_failures: + message.append("---") + + if self.payload.flaky_failures: + flaky_failures = self.payload.flaky_failures + if flaky_failures: + flake_content = [ + f"{messagify_flake(failure)}" for failure in flaky_failures + ] + + flaky_section = wrap_in_details( + f"View the full list of {len(flaky_failures)} :snowflake: flaky tests", + "\n".join(flake_content), + ) + + message.append(flaky_section) + + message.append(generate_view_test_analytics_line(self.commit)) + return "\n".join(message) + + def error_comment(self): + """ + This is no longer used in the new TA finisher task, but this is what used to display the error comment + """ + message: str + + if self.error is None: + message = ":x: We are unable to process any of the uploaded JUnit XML files. Please ensure your files are in the right format." + else: + message = specific_error_message(self.error) + + pull = self.get_pull() + if pull is None: + return False, "no_pull" + + sent_to_provider = self.send_to_provider(pull, message) + + if sent_to_provider is False: + return (False, "torngit_error") + + return (True, "comment_posted") + + def upgrade_comment(self): + pull = self.get_pull() + if pull is None: + return False, "no_pull" + + db_pull = pull.database_pull + provider_pull = pull.provider_pull + if provider_pull is None: + return False, "missing_provider_pull" + + link = get_members_url(db_pull) + + author_username = provider_pull["author"].get("username") + + if not requires_license(): + message = "\n".join( + [ + f"The author of this PR, {author_username}, is not an activated member of this organization on Codecov.", + f"Please [activate this user on Codecov]({link}) to display this PR comment.", + "Coverage data is still being uploaded to Codecov.io for purposes of overall coverage calculations.", + "Please don't hesitate to email us at support@codecov.io with any questions.", + ] + ) + else: + message = "\n".join( + [ + f"The author of this PR, {author_username}, is not activated in your Codecov Self-Hosted installation.", + f"Please [activate this user]({link}) to display this PR comment.", + "Coverage data is still being uploaded to Codecov Self-Hosted for the purposes of overall coverage calculations.", + "Please contact your Codecov On-Premises installation administrator with any questions.", + ] + ) + + sent_to_provider = self.send_to_provider(pull, message) + if sent_to_provider == False: + return (False, "torngit_error") + + return (True, "comment_posted") + + +def not_private_and_free_or_team(repo: Repository): + return not ( + repo.private + and repo.owner.plan + in {**FREE_PLAN_REPRESENTATIONS, **TEAM_PLAN_REPRESENTATIONS} + ) + + +def should_do_flaky_detection(repo: Repository, commit_yaml: UserYaml) -> bool: + has_flaky_configured = read_yaml_field( + commit_yaml, ("test_analytics", "flake_detection"), True + ) + feature_enabled = FLAKY_TEST_DETECTION.check_value( + identifier=repo.repoid, default=True + ) + has_valid_plan_repo_or_owner = not_private_and_free_or_team(repo) + return has_flaky_configured and (feature_enabled or has_valid_plan_repo_or_owner) diff --git a/services/tests/test_ta_utils.py b/services/tests/test_ta_utils.py new file mode 100644 index 000000000..ac2c48073 --- /dev/null +++ b/services/tests/test_ta_utils.py @@ -0,0 +1,317 @@ +import mock +import pytest +from shared.torngit.exceptions import TorngitClientError + +from database.models import UploadError +from database.tests.factories import ( + CommitFactory, + OwnerFactory, + RepositoryFactory, + UploadFactory, +) +from helpers.notifier import NotifierResult +from services.ta_utils import ( + FlakeInfo, + TestFailure, + TestResultsNotificationPayload, + TestResultsNotifier, + generate_failure_info, + should_do_flaky_detection, +) +from services.urls import services_short_dict +from services.yaml import UserYaml + + +def mock_repo_service(): + repo_service = mock.Mock( + post_comment=mock.AsyncMock(), + edit_comment=mock.AsyncMock(), + ) + return repo_service + + +def test_send_to_provider(): + tn = TestResultsNotifier(CommitFactory(), None) + tn._pull = mock.Mock() + tn._pull.database_pull.commentid = None + tn._repo_service = mock_repo_service() + m = dict(id=1) + tn._repo_service.post_comment.return_value = m + + res = tn.send_to_provider(tn._pull, "hello world") + + assert res == True + + tn._repo_service.post_comment.assert_called_with( + tn._pull.database_pull.pullid, "hello world" + ) + assert tn._pull.database_pull.commentid == 1 + + +def test_send_to_provider_edit(): + tn = TestResultsNotifier(CommitFactory(), None) + tn._pull = mock.Mock() + tn._pull.database_pull.commentid = 1 + tn._repo_service = mock_repo_service() + m = dict(id=1) + tn._repo_service.edit_comment.return_value = m + + res = tn.send_to_provider(tn._pull, "hello world") + + assert res == True + tn._repo_service.edit_comment.assert_called_with( + tn._pull.database_pull.pullid, 1, "hello world" + ) + + +def test_send_to_provider_fail(): + tn = TestResultsNotifier(CommitFactory(), None) + tn._pull = mock.Mock() + tn._pull.database_pull.commentid = 1 + tn._repo_service = mock_repo_service() + tn._repo_service.edit_comment.side_effect = TorngitClientError + + res = tn.send_to_provider(tn._pull, "hello world") + + assert res == False + + +def test_generate_failure_info(): + fail = TestFailure( + display_name="testname", + failure_message="hello world", + duration_seconds=1.0, + build_url="https://example.com/build_url", + ) + + res = generate_failure_info(fail) + + assert ( + res + == """ +```python +hello world +``` + +[View](https://example.com/build_url) the CI Build""" + ) + + +def test_build_message(): + fail = TestFailure( + display_name="testname", + failure_message="hello world", + duration_seconds=1.0, + build_url="https://example.com/build_url", + ) + payload = TestResultsNotificationPayload(1, 2, 3, regular_failures=[fail]) + commit = CommitFactory(branch="thing/thing") + tn = TestResultsNotifier(commit, None, None, None, payload) + res = tn.build_message() + + assert ( + res + == f"""### :x: 1 Tests Failed: +| Tests completed | Failed | Passed | Skipped | +|---|---|---|---| +| 3 | 1 | 2 | 3 | +
View the top 1 failed tests by shortest run time + +> +> ```python +> testname +> ``` +> +>
Stack Traces | 1s run time +> +> > +> > ```python +> > hello world +> > ``` +> > +> > [View](https://example.com/build_url) the CI Build +> +>
+ +
+ +To view more test analytics, go to the [Test Analytics Dashboard](https://app.codecov.io/{services_short_dict.get(commit.repository.service)}/{commit.repository.owner.username}/{commit.repository.name}/tests/thing%2Fthing) +:loudspeaker: Thoughts on this report? [Let us know!](https://github.com/codecov/feedback/issues/304)""" + ) + + +def test_build_message_with_flake(): + fail = TestFailure( + display_name="testname", + failure_message="hello world", + duration_seconds=1.0, + build_url="https://example.com/build_url", + flake_info=FlakeInfo(1, 3), + ) + payload = TestResultsNotificationPayload(1, 2, 3, flaky_failures=[fail]) + commit = CommitFactory(branch="test_branch") + tn = TestResultsNotifier(commit, None, None, None, payload) + res = tn.build_message() + + assert ( + res + == f"""### :x: 1 Tests Failed: +| Tests completed | Failed | Passed | Skipped | +|---|---|---|---| +| 3 | 1 | 2 | 3 | +
View the full list of 1 :snowflake: flaky tests + +> +> ```python +> testname +> ``` +> +> **Flake rate in main:** 33.33% (Passed 2 times, Failed 1 times) +>
Stack Traces | 1s run time +> +> > +> > ```python +> > hello world +> > ``` +> > +> > [View](https://example.com/build_url) the CI Build +> +>
+ +
+ +To view more test analytics, go to the [Test Analytics Dashboard](https://app.codecov.io/{services_short_dict.get(commit.repository.service)}/{commit.repository.owner.username}/{commit.repository.name}/tests/{commit.branch}) +:loudspeaker: Thoughts on this report? [Let us know!](https://github.com/codecov/feedback/issues/304)""" + ) + + +def test_notify(mocker): + mocker.patch("helpers.notifier.get_repo_provider_service", return_value=mock.Mock()) + mocker.patch( + "helpers.notifier.fetch_and_update_pull_request_information_from_commit", + return_value=mock.Mock(), + ) + tn = TestResultsNotifier(CommitFactory(), None) + tn.build_message = mock.Mock() + tn.send_to_provider = mock.Mock() + + notification_result = tn.notify() + + assert notification_result == NotifierResult.COMMENT_POSTED + + +def test_notify_fail_torngit_error( + mocker, +): + mocker.patch("helpers.notifier.get_repo_provider_service", return_value=mock.Mock()) + mocker.patch( + "helpers.notifier.fetch_and_update_pull_request_information_from_commit", + return_value=mock.Mock(), + ) + tn = TestResultsNotifier(CommitFactory(), None) + tn.build_message = mock.Mock() + tn.send_to_provider = mock.Mock(return_value=False) + + notification_result = tn.notify() + + assert notification_result == NotifierResult.TORNGIT_ERROR + + +def test_notify_fail_no_pull( + mocker, +): + mocker.patch("helpers.notifier.get_repo_provider_service", return_value=mock.Mock()) + mocker.patch( + "helpers.notifier.fetch_and_update_pull_request_information_from_commit", + return_value=None, + ) + tn = TestResultsNotifier(CommitFactory(), None) + tn.build_message = mock.Mock() + tn.send_to_provider = mock.Mock(return_value=False) + + notification_result = tn.notify() + assert notification_result == NotifierResult.NO_PULL + + +@pytest.mark.parametrize( + "config,feature_flag,private,plan,ex_result", + [ + (False, True, False, "users-inappm", False), + (True, True, True, "users-basic", True), + (True, False, False, "users-basic", True), + (True, False, True, "users-basic", False), + (True, False, False, "users-inappm", True), + (True, False, True, "users-inappm", True), + ], +) +def test_should_do_flake_detection( + dbsession, mocker, config, feature_flag, private, plan, ex_result +): + owner = OwnerFactory(plan=plan) + repo = RepositoryFactory(private=private, owner=owner) + dbsession.add(repo) + dbsession.flush() + + mocked_feature = mocker.patch("services.ta_utils.FLAKY_TEST_DETECTION") + mocked_feature.check_value.return_value = feature_flag + + yaml = {"test_analytics": {"flake_detection": config}} + + result = should_do_flaky_detection(repo, UserYaml.from_dict(yaml)) + + assert result == ex_result + + +def test_specific_error_message(mocker): + mock_repo_service = mock.AsyncMock() + mocker.patch( + "helpers.notifier.get_repo_provider_service", return_value=mock_repo_service + ) + mocker.patch( + "helpers.notifier.fetch_and_update_pull_request_information_from_commit", + return_value=mock.AsyncMock(), + ) + + upload = UploadFactory() + error = UploadError( + report_upload=upload, + error_code="Unsupported file format", + error_params={ + "error_message": "Error parsing JUnit XML in test.xml at 4:32: ParserError: No name found" + }, + ) + tn = TestResultsNotifier(CommitFactory(), None, error=error) + result = tn.error_comment() + expected = """### :x: Unsupported file format + +> Upload processing failed due to unsupported file format. Please review the parser error message: +> `Error parsing JUnit XML in test.xml at 4:32: ParserError: No name found` +> For more help, visit our [troubleshooting guide](https://docs.codecov.com/docs/test-analytics#troubleshooting). +""" + + assert result == (True, "comment_posted") + print(mock_repo_service.mock_calls) + mock_repo_service.edit_comment.assert_called_with( + tn._pull.database_pull.pullid, tn._pull.database_pull.commentid, expected + ) + + +def test_specific_error_message_no_error(mocker): + mock_repo_service = mock.AsyncMock() + mocker.patch( + "helpers.notifier.get_repo_provider_service", return_value=mock_repo_service + ) + mocker.patch( + "helpers.notifier.fetch_and_update_pull_request_information_from_commit", + return_value=mock.AsyncMock(), + ) + + tn = TestResultsNotifier(CommitFactory(), None) + result = tn.error_comment() + expected = """:x: We are unable to process any of the uploaded JUnit XML files. Please ensure your files are in the right format.""" + + assert result == (True, "comment_posted") + print(mock_repo_service.mock_calls) + mock_repo_service.edit_comment.assert_called_with( + tn._pull.database_pull.pullid, tn._pull.database_pull.commentid, expected + ) diff --git a/ta_storage/base.py b/ta_storage/base.py index 690e3fefa..fa9b347bd 100644 --- a/ta_storage/base.py +++ b/ta_storage/base.py @@ -1,23 +1,68 @@ from __future__ import annotations from abc import ABC, abstractmethod +from typing import Generic, TypedDict, TypeVar import test_results_parser +from shared.django_apps.reports.models import ReportSession -from database.models.reports import Upload +from services.ta_utils import FlakeInfo -class TADriver(ABC): +class PRCommentAggResult(TypedDict): + commit_sha: str + passed_ct: int + failed_ct: int + skipped_ct: int + flaky_failed_ct: int + + +T = TypeVar("T") + + +class PRCommentFailResult(TypedDict, Generic[T]): + id: T + computed_name: str + failure_message: str | None + duration_seconds: float + upload_id: int + + +class TADriver(ABC, Generic[T]): + def __init__(self, repo_id: int) -> None: + self.repo_id = repo_id + @abstractmethod def write_testruns( self, - timestamp: int, - repo_id: int, + timestamp: int | None, commit_sha: str, branch_name: str, - upload: Upload, + upload_id: int, + flag_names: list[str], framework: str | None, testruns: list[test_results_parser.Testrun], - flaky_test_set: set[str], - ): + ) -> None: + pass + + @abstractmethod + def write_flakes(self, uploads: list[ReportSession]) -> None: + pass + + @abstractmethod + def cache_analytics(self, buckets: list[str], branch: str) -> None: + pass + + @abstractmethod + def pr_comment_agg(self, commit_sha: str) -> PRCommentAggResult: + pass + + @abstractmethod + def pr_comment_fail(self, commit_sha: str) -> list[PRCommentFailResult[T]]: + pass + + @abstractmethod + def get_repo_flakes( + self, test_ids: tuple[T, ...] | None = None + ) -> dict[T, FlakeInfo]: pass diff --git a/ta_storage/bq.py b/ta_storage/bq.py index 7d9658c75..a964c0940 100644 --- a/ta_storage/bq.py +++ b/ta_storage/bq.py @@ -1,18 +1,42 @@ from __future__ import annotations +import datetime as dt from datetime import datetime -from typing import Literal, cast +from functools import cached_property +from typing import Literal, TypedDict, cast +import polars as pl +import shared.storage import test_results_parser -from google.cloud.bigquery import ArrayQueryParameter, ScalarQueryParameter +from django.db import transaction +from google.cloud.bigquery import ( + ArrayQueryParameter, + ScalarQueryParameter, + StructQueryParameter, + StructQueryParameterType, +) from shared.config import get_config +from shared.django_apps.reports.models import ReportSession +from shared.django_apps.test_analytics.models import Flake import generated_proto.testrun.ta_testrun_pb2 as ta_testrun_pb2 -from database.models.reports import Upload from services.bigquery import get_bigquery_service -from ta_storage.base import TADriver +from services.ta_utils import FlakeInfo +from ta_storage.base import ( + PRCommentAggResult, + PRCommentFailResult, + TADriver, +) from ta_storage.utils import calc_flags_hash, calc_test_id + +def get_live_flakes(repo_id: int) -> list[Flake]: + return [ + flake + for flake in Flake.objects.filter(repoid=repo_id, end_date__isnull=True).all() + ] + + RANKED_DATA = """ ranked_data AS ( SELECT @@ -68,7 +92,10 @@ SELECT computed_name, failure_message, - flags + test_id, + flags_hash, + duration_seconds, + upload_id FROM latest_instances WHERE @@ -77,10 +104,11 @@ TESTRUNS_FOR_UPLOAD = """ SELECT - DATE_BUCKET(timestamp, INTERVAL 1 DAY) AS date, + timestamp, test_id, outcome, branch_name, + flags_hash, FROM `{PROJECT_ID}.{DATASET_NAME}.{TESTRUN_TABLE_NAME}` WHERE @@ -88,7 +116,12 @@ AND ( outcome = 1 OR outcome = 3 - OR test_id IN UNNEST(@test_ids) + OR EXISTS ( + SELECT 1 + FROM UNNEST(@flake_ids) AS flake_id + WHERE flake_id.candidate_test_id = test_id + AND IF(flake_id.candidate_flags_hash IS NULL, flags_hash IS NULL, flake_id.candidate_flags_hash = flags_hash) + ) ) """ @@ -135,6 +168,29 @@ """ +class TestrunsForUploadResult(TypedDict): + timestamp: int + test_id: bytes + flags_hash: bytes + outcome: Literal["passed", "failed", "skipped", "flaky_failed"] + branch_name: str + + +def get_storage_key( + repo_id: int, branch: str | None, interval_start: int, interval_end: int | None +) -> str: + interval_section = ( + f"{interval_start}" + if interval_end is None + else f"{interval_start}_{interval_end}" + ) + + if branch: + return f"ta_rollups/{repo_id}/{branch}/{interval_section}" + else: + return f"ta_rollups/{repo_id}/{interval_section}" + + def outcome_to_enum( outcome: Literal["pass", "skip", "failure", "error"], ) -> ta_testrun_pb2.TestRun.Outcome: @@ -149,9 +205,9 @@ def outcome_to_enum( raise ValueError(f"Invalid outcome: {outcome}") -class BQDriver(TADriver): - def __init__(self, flaky_test_set: set[bytes] | None = None): - self.flaky_test_set = flaky_test_set or {} +class BQDriver(TADriver[tuple[bytes, bytes | None]]): + def __init__(self, repo_id: int) -> None: + super().__init__(repo_id) self.bq_service = get_bigquery_service() self.project_id: str = cast( @@ -170,34 +226,44 @@ def __init__(self, flaky_test_set: set[bytes] | None = None): ), ) + def get_live_flakes(self, repo_id: int) -> list[Flake]: + return [ + flake + for flake in Flake.objects.filter( + repoid=repo_id, end_date__isnull=True + ).all() + ] + def write_testruns( self, timestamp: int | None, - repo_id: int, commit_sha: str, branch_name: str, - upload: Upload, + upload_id: int, + flag_names: list[str], framework: str | None, testruns: list[test_results_parser.Testrun], ): if timestamp is None: timestamp = int(datetime.now().timestamp() * 1000000) - flag_names: list[str] = upload.flag_names testruns_pb: list[bytes] = [] flags_hash = calc_flags_hash(flag_names) + flakes: list[Flake] = list(self.flake_dict.values()) + flake_set = {(flake.test_id, flake.flags_id) for flake in flakes} + for t in testruns: test_id = calc_test_id(t["name"], t["classname"], t["testsuite"]) - if test_id in self.flaky_test_set and t["outcome"] == "failure": + if (test_id, flags_hash) in flake_set and t["outcome"] == "failure": outcome = ta_testrun_pb2.TestRun.Outcome.FLAKY_FAILED else: outcome = outcome_to_enum(t["outcome"]) test_run = ta_testrun_pb2.TestRun( timestamp=timestamp, - repoid=repo_id, + repoid=self.repo_id, commit_sha=commit_sha, framework=framework, branch_name=branch_name, @@ -205,12 +271,13 @@ def write_testruns( classname=t["classname"], name=t["name"], testsuite=t["testsuite"], - computed_name=t["computed_name"], + computed_name=t["computed_name"] + or f"{t['testsuite']}.{t['classname']}.{t['name']}", outcome=outcome, failure_message=t["failure_message"], duration_seconds=t["duration"], filename=t["filename"], - upload_id=upload.id_, + upload_id=upload_id, flags_hash=flags_hash, test_id=test_id, ) @@ -220,11 +287,7 @@ def write_testruns( self.dataset_name, self.testrun_table_name, ta_testrun_pb2, testruns_pb ) - def pr_comment_agg( - self, - repoid: int, - commit_sha: str, - ): + def pr_comment_agg(self, commit_sha: str) -> PRCommentAggResult: query = f""" WITH {RANKED_DATA.format( @@ -235,15 +298,27 @@ def pr_comment_agg( {LATEST_INSTANCES} {PR_COMMENT_AGG} """ - return self.bq_service.query( + query_result = self.bq_service.query( query, [ - ScalarQueryParameter("repoid", "INT64", repoid), + ScalarQueryParameter("repoid", "INT64", self.repo_id), ScalarQueryParameter("commit_sha", "STRING", commit_sha), ], ) - def pr_comment_fail(self, repoid: int, commit_sha: str): + result = query_result[0] + + return { + "commit_sha": result["commit_sha"], + "passed_ct": result["passed"], + "failed_ct": result["failed"], + "skipped_ct": result["skipped"], + "flaky_failed_ct": result["flaky_failed"], + } + + def pr_comment_fail( + self, commit_sha: str + ) -> list[PRCommentFailResult[tuple[bytes, bytes | None]]]: query = f""" WITH {RANKED_DATA.format( @@ -254,15 +329,28 @@ def pr_comment_fail(self, repoid: int, commit_sha: str): {LATEST_INSTANCES} {PR_COMMENT_FAIL} """ - return self.bq_service.query( + query_result = self.bq_service.query( query, [ - ScalarQueryParameter("repoid", "INT64", repoid), + ScalarQueryParameter("repoid", "INT64", self.repo_id), ScalarQueryParameter("commit_sha", "STRING", commit_sha), ], ) - def testruns_for_upload(self, upload_id: int, test_ids: list[bytes]): + return [ + { + "computed_name": result["computed_name"], + "failure_message": result["failure_message"], + "id": (result["test_id"], result["flags_hash"]), + "duration_seconds": result["duration_seconds"], + "upload_id": result["upload_id"], + } + for result in query_result + ] + + def testruns_for_upload( + self, upload_id: int, flake_set: set[tuple[bytes, bytes | None]] + ) -> list[TestrunsForUploadResult]: query = f""" {TESTRUNS_FOR_UPLOAD.format( PROJECT_ID=self.project_id, @@ -270,14 +358,119 @@ def testruns_for_upload(self, upload_id: int, test_ids: list[bytes]): TESTRUN_TABLE_NAME=self.testrun_table_name, )} """ - return self.bq_service.query( + # IMPORTANT: the names of these parameters actually matters or + # else BQ won't be able to disambiguate the test_id of the object + # up for comparison and the field in the struct + # if the name of the fields in the structs are equal to the name of the fields + # in the table, it will match every row + ids = [ + StructQueryParameter( + None, + ScalarQueryParameter("candidate_test_id", "BYTES", test_id), + ScalarQueryParameter("candidate_flags_hash", "BYTES", flags_hash), + ) + for (test_id, flags_hash) in flake_set + ] + + query_result = self.bq_service.query( query, [ ScalarQueryParameter("upload_id", "INT64", upload_id), - ArrayQueryParameter("test_ids", "BYTES", test_ids), + ArrayQueryParameter( + "flake_ids", + StructQueryParameterType( + ScalarQueryParameter("candidate_test_id", "BYTES", None), + ScalarQueryParameter("candidate_flags_hash", "BYTES", None), + ), + ids, + ), ], ) + return [ + { + "branch_name": result["branch_name"], + "timestamp": result["timestamp"], + "outcome": result["outcome"], + "test_id": result["test_id"], + "flags_hash": result["flags_hash"], + } + for result in query_result + ] + + @cached_property + def flake_dict(self) -> dict[tuple[bytes, bytes | None], Flake]: + return { + ( + bytes(flake.test_id), + bytes(flake.flags_id) if flake.flags_id else None, + ): flake + for flake in self.get_live_flakes(self.repo_id) + } + + def write_flakes(self, uploads: list[ReportSession]) -> None: + # get relevant flakes test ids: + flakes = list(self.flake_dict.values()) + + flake_dict = { + ( + bytes(flake.test_id), + bytes(flake.flags_id) if flake.flags_id else None, + ): flake + for flake in flakes + } + + for upload in uploads: + testruns = self.testruns_for_upload(upload.id, set(flake_dict.keys())) + for testrun in testruns: + if flake := flake_dict.get((testrun["test_id"], testrun["flags_hash"])): + match testrun["outcome"]: + case ( + ta_testrun_pb2.TestRun.Outcome.FAILED + | ta_testrun_pb2.TestRun.Outcome.FLAKY_FAILED + ): + flake.fail_count += 1 + flake.count += 1 + flake.recent_passes_count = 0 + case ta_testrun_pb2.TestRun.Outcome.PASSED: + flake.recent_passes_count += 1 + flake.count += 1 + + if flake.recent_passes_count == 30: + flake.end_date = datetime.now() + case _: + pass + + flake.save() + else: + match testrun["outcome"]: + case ( + ta_testrun_pb2.TestRun.Outcome.FAILED + | ta_testrun_pb2.TestRun.Outcome.FLAKY_FAILED + ): + flake = Flake.objects.create( + repoid=self.repo_id, + test_id=testrun["test_id"], + fail_count=1, + count=1, + recent_passes_count=0, + start_date=datetime.fromtimestamp( + testrun["timestamp"] / 1000000 + ), + end_date=None, + ) + flake.save() + + flake_dict[(testrun["test_id"], testrun["flags_hash"])] = ( + flake + ) + case _: + pass + + upload.state = "flake_processed" + upload.save() + transaction.commit() + def analytics( self, repoid: int, @@ -316,3 +509,90 @@ def analytics( params.append(ScalarQueryParameter("branch", "STRING", branch)) return self.bq_service.query(query, params) + + def cache_analytics(self, buckets: list[str], branch: str | None) -> None: + storage_service = shared.storage.get_appropriate_storage_service(self.repo_id) + + for interval_start, interval_end in [ + # NOTE: working with calendar days and intervals, + # `(CURRENT_DATE - INTERVAL '1 days')` means *yesterday*, + # and `2..1` matches *the day before yesterday*. + (1, None), + (2, 1), + (7, None), + (14, 7), + (30, None), + (60, 30), + ]: + analytics_results = self.analytics( + self.repo_id, + interval_start=interval_start, + interval_end=interval_end, + branch=branch, + ) + + df = pl.DataFrame( + analytics_results, + [ + "name", + "classname", + "testsuite", + "computed_name", + ("flags", pl.List(pl.String)), + "test_id", + ("updated_at", pl.Datetime(time_zone=dt.UTC)), + "avg_duration", + "fail_count", + "flaky_fail_count", + "pass_count", + "skip_count", + "commits_where_fail", + "last_duration", + ], + orient="row", + ) + + serialized_table = df.write_ipc(None) + serialized_table.seek(0) # avoids Stream must be at beginning errors + + storage_key = get_storage_key( + self.repo_id, branch, interval_start, interval_end + ) + for bucket in buckets: + storage_service.write_file( + bucket, + storage_key, + serialized_table, + ) + + def get_repo_flakes( + self, test_ids: tuple[tuple[bytes, bytes | None], ...] | None = None + ) -> dict[tuple[bytes, bytes | None], FlakeInfo]: + if test_ids: + return { + ( + bytes(flake.test_id), + bytes(flake.flags_id) if flake.flags_id else None, + ): FlakeInfo( + failed=flake.fail_count, + count=flake.count, + ) + for flake in Flake.objects.raw( + "SELECT * FROM flake WHERE repoid = %s AND (test_id, flags_id) IN %s AND end_date IS NULL AND count != recent_passes_count + fail_count", + [self.repo_id, test_ids], + ).all() + } + else: + return { + ( + bytes(flake.test_id), + bytes(flake.flags_id) if flake.flags_id else None, + ): FlakeInfo( + failed=flake.fail_count, + count=flake.count, + ) + for flake in Flake.objects.raw( + "SELECT * FROM flake WHERE repoid = %s AND end_date IS NULL AND count != recent_passes_count + fail_count", + [self.repo_id], + ).all() + } diff --git a/ta_storage/pg.py b/ta_storage/pg.py index 762efd8bd..f6cca349d 100644 --- a/ta_storage/pg.py +++ b/ta_storage/pg.py @@ -1,22 +1,144 @@ from __future__ import annotations +import datetime as dt from datetime import date, datetime from typing import Any, Literal, TypedDict +import polars as pl +import shared.storage import test_results_parser +from django.db import connections +from django.db import transaction as django_transaction +from django.db.models import Q +from shared.config import get_config +from shared.django_apps.reports.models import ( + CommitReport as DjangoCommitReport, +) +from shared.django_apps.reports.models import ( + DailyTestRollup as DjangoDailyTestRollup, +) +from shared.django_apps.reports.models import ( + Flake as DjangoFlake, +) +from shared.django_apps.reports.models import ( + ReportSession as DjangoReportSession, +) +from shared.django_apps.reports.models import ( + TestInstance as DjangoTestInstance, +) from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session from database.models import ( DailyTestRollup, + Flake, RepositoryFlag, Test, TestFlagBridge, TestInstance, - Upload, ) -from services.test_results import generate_flags_hash, generate_test_id -from ta_storage.base import TADriver +from services.ta_utils import ( + FlakeInfo, + generate_flags_hash, + generate_test_id, + get_test_summary_for_commit, + latest_failures_for_commit, +) +from ta_storage.base import ( + PRCommentAggResult, + PRCommentFailResult, + TADriver, +) + +# Reminder: `a BETWEEN x AND y` is equivalent to `a >= x AND a <= y` +# Since we are working with calendar days, using a range of `0..0` gives us *today*, +# and a range of `1..1` gives use *yesterday*. +BASE_SUBQUERY = """ +SELECT * +FROM reports_dailytestrollups +WHERE repoid = %(repoid)s + AND branch = %(branch)s + AND date BETWEEN + (CURRENT_DATE - INTERVAL %(interval_start)s) AND + (CURRENT_DATE - INTERVAL %(interval_end)s) +""" + +TEST_AGGREGATION_SUBQUERY = """ +SELECT test_id, + CASE + WHEN SUM(pass_count) + SUM(fail_count) = 0 THEN 0 + ELSE SUM(fail_count)::float / (SUM(pass_count) + SUM(fail_count)) + END AS failure_rate, + CASE + WHEN SUM(pass_count) + SUM(fail_count) = 0 THEN 0 + ELSE SUM(flaky_fail_count)::float / (SUM(pass_count) + SUM(fail_count)) + END AS flake_rate, + MAX(latest_run) AS updated_at, + AVG(avg_duration_seconds) AS avg_duration, + SUM(fail_count) AS total_fail_count, + SUM(flaky_fail_count) AS total_flaky_fail_count, + SUM(pass_count) AS total_pass_count, + SUM(skip_count) AS total_skip_count +FROM base_cte +GROUP BY test_id +""" + +COMMITS_FAILED_SUBQUERY = """ +SELECT test_id, + array_length((array_agg(DISTINCT unnested_cwf)), 1) AS failed_commits_count +FROM + (SELECT test_id, + commits_where_fail AS cwf + FROM base_cte + WHERE array_length(commits_where_fail, 1) > 0) AS tests_with_commits_that_failed, + unnest(cwf) AS unnested_cwf +GROUP BY test_id +""" + +LAST_DURATION_SUBQUERY = """ +SELECT base_cte.test_id, + last_duration_seconds +FROM base_cte +JOIN + (SELECT test_id, + max(created_at) AS created_at + FROM base_cte + GROUP BY test_id) AS latest_rollups ON base_cte.created_at = latest_rollups.created_at +AND base_cte.test_id = latest_rollups.test_id +""" + +TEST_FLAGS_SUBQUERY = """ +SELECT test_id, + array_agg(DISTINCT flag_name) AS flags +FROM reports_test_results_flag_bridge tfb +JOIN reports_test rt ON rt.id = tfb.test_id +JOIN reports_repositoryflag rr ON tfb.flag_id = rr.id +WHERE rt.repoid = %(repoid)s +GROUP BY test_id +""" + +ROLLUP_QUERY = f""" +WITH + base_cte AS ({BASE_SUBQUERY}), + failure_rate_cte AS ({TEST_AGGREGATION_SUBQUERY}), + commits_where_fail_cte AS ({COMMITS_FAILED_SUBQUERY}), + last_duration_cte AS ({LAST_DURATION_SUBQUERY}), + flags_cte AS ({TEST_FLAGS_SUBQUERY}) + +SELECT COALESCE(rt.computed_name, rt.name) AS name, + rt.testsuite, + flags_cte.flags, + results.* +FROM + (SELECT failure_rate_cte.*, + coalesce(commits_where_fail_cte.failed_commits_count, 0) AS commits_where_fail, + last_duration_cte.last_duration_seconds AS last_duration + FROM failure_rate_cte + FULL OUTER JOIN commits_where_fail_cte USING (test_id) + FULL OUTER JOIN last_duration_cte USING (test_id)) AS results +JOIN reports_test rt ON results.test_id = rt.id +LEFT JOIN flags_cte USING (test_id) +""" class DailyTotals(TypedDict): @@ -54,7 +176,8 @@ def modify_structures( test_flag_bridge_data: list[dict], daily_totals: dict[str, DailyTotals], testrun: test_results_parser.Testrun, - upload: Upload, + upload_id: int, + flag_names: list[str], repoid: int, branch: str | None, commit_sha: str, @@ -62,7 +185,7 @@ def modify_structures( flaky_test_set: set[str], framework: str | None, ): - flags_hash = generate_flags_hash(upload.flag_names) + flags_hash = generate_flags_hash(flag_names) test_name = f"{testrun['classname']}\x1f{testrun['name']}" test_id = generate_test_id( @@ -78,7 +201,7 @@ def modify_structures( tests_to_write[test_id] = test test_instance = generate_test_instance_dict( - test_id, upload, testrun, commit_sha, branch, repoid + test_id, upload_id, testrun, commit_sha, branch, repoid ) test_instances_to_write.append(test_instance) @@ -129,7 +252,7 @@ def generate_test_dict( def generate_test_instance_dict( test_id: str, - upload: Upload, + upload_id: int, testrun: test_results_parser.Testrun, commit_sha: str, branch: str | None, @@ -137,7 +260,7 @@ def generate_test_instance_dict( ) -> dict[str, Any]: return { "test_id": test_id, - "upload_id": upload.id, + "upload_id": upload_id, "duration_seconds": testrun["duration"], "outcome": testrun["outcome"], "failure_message": testrun["failure_message"], @@ -282,27 +405,220 @@ def save_test_instances(db_session: Session, test_instance_data: list[dict]): db_session.commit() -class PGDriver(TADriver): - def __init__(self, db_session: Session, flaky_test_set: set[str]): +FLAKE_EXPIRY_COUNT = 30 + + +def process_flake_for_repo_commit( + repo_id: int, + commit_id: str, +): + uploads = DjangoReportSession.objects.filter( + report__report_type=DjangoCommitReport.ReportType.TEST_RESULTS.value, + report__commit__repository__repoid=repo_id, + report__commit__commitid=commit_id, + state__in=["processed", "v2_finished"], + ).all() + + process_flakes_for_uploads(repo_id, [upload for upload in uploads]) + + return {"successful": True} + + +def process_flakes_for_uploads(repo_id: int, uploads: list[DjangoReportSession]): + curr_flakes = fetch_curr_flakes(repo_id) + new_flakes: dict[str, DjangoFlake] = dict() + + rollups_to_update: list[DjangoDailyTestRollup] = [] + + flaky_tests = list(curr_flakes.keys()) + + for upload in uploads: + test_instances = get_test_instances(upload, flaky_tests) + for test_instance in test_instances: + if test_instance.outcome == DjangoTestInstance.Outcome.PASS.value: + flake = new_flakes.get(test_instance.test_id) or curr_flakes.get( + test_instance.test_id + ) + if flake is not None: + update_flake(flake, test_instance) + elif test_instance.outcome in ( + DjangoTestInstance.Outcome.FAILURE.value, + DjangoTestInstance.Outcome.ERROR.value, + ): + flake = new_flakes.get(test_instance.test_id) or curr_flakes.get( + test_instance.test_id + ) + if flake: + update_flake(flake, test_instance) + else: + flake, rollup = create_flake(test_instance, repo_id) + + new_flakes[test_instance.test_id] = flake + + if rollup: + rollups_to_update.append(rollup) + + if rollups_to_update: + DjangoDailyTestRollup.objects.bulk_update( + rollups_to_update, + ["flaky_fail_count"], + ) + + merge_flake_dict = {} + + if new_flakes: + flakes_to_merge = DjangoFlake.objects.bulk_create(new_flakes.values()) + merge_flake_dict: dict[str, DjangoFlake] = { + flake.test_id: flake for flake in flakes_to_merge + } + + DjangoFlake.objects.bulk_update( + curr_flakes.values(), + [ + "count", + "fail_count", + "recent_passes_count", + "end_date", + ], + ) + + curr_flakes = {**merge_flake_dict, **curr_flakes} + + new_flakes.clear() + + upload.state = "flake_processed" + upload.save() + django_transaction.commit() + + +def get_test_instances( + upload: DjangoReportSession, + flaky_tests: list[str], +) -> list[DjangoTestInstance]: + # get test instances on this upload that either: + # - failed + # - passed but belong to an already flaky test + + upload_filter = Q(upload_id=upload.id) + test_failed_filter = Q(outcome=DjangoTestInstance.Outcome.ERROR.value) | Q( + outcome=DjangoTestInstance.Outcome.FAILURE.value + ) + test_passed_but_flaky_filter = Q(outcome=DjangoTestInstance.Outcome.PASS.value) & Q( + test_id__in=flaky_tests + ) + test_instances = list( + DjangoTestInstance.objects.filter( + upload_filter & (test_failed_filter | test_passed_but_flaky_filter) + ) + .select_related("test") + .all() + ) + return test_instances + + +def fetch_curr_flakes(repo_id: int) -> dict[str, DjangoFlake]: + flakes = DjangoFlake.objects.filter( + repository_id=repo_id, end_date__isnull=True + ).all() + return {flake.test_id: flake for flake in flakes} + + +def create_flake( + test_instance: DjangoTestInstance, + repo_id: int, +) -> tuple[DjangoFlake, DjangoDailyTestRollup | None]: + # retroactively mark newly caught flake as flaky failure in its rollup + rollup = DjangoDailyTestRollup.objects.filter( + repoid=repo_id, + date=test_instance.created_at.date(), + branch=test_instance.branch, + test_id=test_instance.test_id, + ).first() + + if rollup: + rollup.flaky_fail_count += 1 + + f = DjangoFlake( + repository_id=repo_id, + test=test_instance.test, + reduced_error=None, + count=1, + fail_count=1, + start_date=test_instance.created_at, + recent_passes_count=0, + ) + + return f, rollup + + +def update_flake( + flake: DjangoFlake, + test_instance: DjangoTestInstance, +) -> None: + flake.count += 1 + + match test_instance.outcome: + case DjangoTestInstance.Outcome.PASS.value: + flake.recent_passes_count += 1 + if flake.recent_passes_count == FLAKE_EXPIRY_COUNT: + flake.end_date = test_instance.created_at + case ( + DjangoTestInstance.Outcome.FAILURE.value + | DjangoTestInstance.Outcome.ERROR.value + ): + flake.fail_count += 1 + flake.recent_passes_count = 0 + case _: + pass + + +def get_storage_key( + repo_id: int, branch: str | None, interval_start: int, interval_end: int | None +) -> str: + interval_section = ( + f"{interval_start}" + if interval_end is None + else f"{interval_start}_{interval_end}" + ) + if branch: + return f"ta_rollups/{repo_id}/{branch}/{interval_section}" + else: + return f"ta_rollups/{repo_id}/{interval_section}" + + +class PGDriver(TADriver[str]): + def __init__( + self, + repo_id: int, + db_session: Session | None = None, + flaky_test_set: set[str] | None = None, + ): + super().__init__(repo_id) self.db_session = db_session self.flaky_test_set = flaky_test_set def write_testruns( self, timestamp: int | None, - repo_id: int, commit_sha: str, branch_name: str, - upload: Upload, + upload_id: int, + flag_names: list[str], framework: str | None, testruns: list[test_results_parser.Testrun], ): + if self.db_session is None: + raise ValueError("DB session is required") + + if self.flaky_test_set is None: + self.flaky_test_set = set(self.get_repo_flakes().keys()) + tests_to_write: dict[str, dict[str, Any]] = {} test_instances_to_write: list[dict[str, Any]] = [] daily_totals: dict[str, DailyTotals] = dict() test_flag_bridge_data: list[dict] = [] - repo_flag_ids = get_repo_flag_ids(self.db_session, repo_id, upload.flag_names) + repo_flag_ids = get_repo_flag_ids(self.db_session, self.repo_id, flag_names) for testrun in testruns: modify_structures( @@ -311,8 +627,9 @@ def write_testruns( test_flag_bridge_data, daily_totals, testrun, - upload, - repo_id, + upload_id, + flag_names, + self.repo_id, branch_name, commit_sha, repo_flag_ids, @@ -331,3 +648,140 @@ def write_testruns( if len(test_instances_to_write) > 0: save_test_instances(self.db_session, test_instances_to_write) + + def cache_analytics(self, buckets: list[str], branch: str | None) -> None: + storage_service = shared.storage.get_appropriate_storage_service(self.repo_id) + + if get_config("setup", "database", "read_replica_enabled", default=False): + connection = connections["default_read"] + else: + connection = connections["default"] + + with connection.cursor() as cursor: + for interval_start, interval_end in [ + # NOTE: working with calendar days and intervals, + # `(CURRENT_DATE - INTERVAL '1 days')` means *yesterday*, + # and `2..1` matches *the day before yesterday*. + (1, None), + (2, 1), + (7, None), + (14, 7), + (30, None), + (60, 30), + ]: + query_params = { + "repoid": self.repo_id, + "branch": branch, + "interval_start": f"{interval_start} days", + # SQL `BETWEEN` syntax is equivalent to `<= end`, with an inclusive end date, + # thats why we do a `+1` here: + "interval_end": f"{interval_end + 1 if interval_end else 0} days", + } + + cursor.execute(ROLLUP_QUERY, query_params) + aggregation_of_test_results = cursor.fetchall() + + df = pl.DataFrame( + aggregation_of_test_results, + [ + "name", + "testsuite", + ("flags", pl.List(pl.String)), + "test_id", + "failure_rate", + "flake_rate", + ("updated_at", pl.Datetime(time_zone=dt.UTC)), + "avg_duration", + "total_fail_count", + "total_flaky_fail_count", + "total_pass_count", + "total_skip_count", + "commits_where_fail", + "last_duration", + ], + orient="row", + ) + + serialized_table = df.write_ipc(None) + serialized_table.seek(0) # avoids Stream must be at beginning errors + + storage_key = ( + f"test_results/rollups/{self.repo_id}/{branch}/{interval_start}" + if interval_end is None + else f"test_results/rollups/{self.repo_id}/{branch}/{interval_start}_{interval_end}" + ) + + for bucket in buckets: + storage_service.write_file(bucket, storage_key, serialized_table) + + def pr_comment_agg(self, commit_sha: str) -> PRCommentAggResult: + if self.db_session is None: + raise ValueError("DB session is required") + + test_summary = get_test_summary_for_commit( + self.db_session, self.repo_id, commit_sha + ) + + return { + "commit_sha": commit_sha, + "passed_ct": test_summary.get("pass", 0), + "failed_ct": test_summary.get("failure", 0) + test_summary.get("error", 0), + "skipped_ct": test_summary.get("skip", 0), + "flaky_failed_ct": 0, + } + + def pr_comment_fail(self, commit_sha: str) -> list[PRCommentFailResult[str]]: + if self.db_session is None: + raise ValueError("DB session is required") + + test_instances = latest_failures_for_commit( + self.db_session, self.repo_id, commit_sha + ) + + return [ + { + "computed_name": instance.test.computed_name + or f"{instance.test.testsuite}.{instance.test.name.replace('\x1f', '.')}", + "failure_message": instance.failure_message, + "id": instance.test.id, + "duration_seconds": instance.duration_seconds, + "upload_id": instance.upload_id, + } + for instance in test_instances + ] + + def write_flakes(self, uploads: list[DjangoReportSession]) -> None: + return process_flakes_for_uploads(self.repo_id, uploads) + + def get_repo_flakes( + self, test_ids: tuple[str, ...] | None = None + ) -> dict[str, FlakeInfo]: + if self.db_session is None: + raise ValueError("DB session is required") + if test_ids: + matching_flakes = list( + self.db_session.query(Flake) + .filter( + Flake.repoid == self.repo_id, + Flake.testid.in_(test_ids), + Flake.end_date.is_(None), + Flake.count != (Flake.recent_passes_count + Flake.fail_count), + ) + .all() + ) + else: + matching_flakes = list( + self.db_session.query(Flake) + .filter( + Flake.repoid == self.repo_id, + Flake.end_date.is_(None), + Flake.count != (Flake.recent_passes_count + Flake.fail_count), + ) + .all() + ) + + flaky_test_ids = { + flake.testid: FlakeInfo(flake.fail_count, flake.count) + for flake in matching_flakes + } + return flaky_test_ids diff --git a/ta_storage/tests/snapshots/bq__analytics__0.txt b/ta_storage/tests/snapshots/bq__analytics__0.txt new file mode 100644 index 000000000..46807bb9e --- /dev/null +++ b/ta_storage/tests/snapshots/bq__analytics__0.txt @@ -0,0 +1,31 @@ + + WITH + +analytics_base AS ( + SELECT * + FROM `test-project.test-dataset.test-table` + WHERE repoid = @repoid + AND timestamp BETWEEN + (CURRENT_DATE - INTERVAL @interval_start) AND + (CURRENT_DATE - INTERVAL @interval_end) +) + + +SELECT + name, + classname, + testsuite, + ANY_VALUE(computed_name) AS computed_name, + COUNT(DISTINCT IF(outcome = 1 OR outcome = 3, commit_sha, NULL)) AS cwf, + AVG(duration_seconds) AS avg_duration, + MAX_BY(duration_seconds, timestamp) AS last_duration, + SUM(IF(outcome = 0, 1, 0)) AS pass_count, + SUM(IF(outcome = 1, 1, 0)) AS fail_count, + SUM(IF(outcome = 2, 1, 0)) AS skip_count, + SUM(IF(outcome = 3, 1, 0)) AS flaky_fail_count, + MAX(timestamp) AS updated_at, + ARRAY_AGG(DISTINCT unnested_flags) AS flags +FROM analytics_base, UNNEST(flags) AS unnested_flags +GROUP BY name, classname, testsuite + + \ No newline at end of file diff --git a/ta_storage/tests/snapshots/bq__analytics_with_branch__0.txt b/ta_storage/tests/snapshots/bq__analytics_with_branch__0.txt new file mode 100644 index 000000000..6785389e7 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__analytics_with_branch__0.txt @@ -0,0 +1,32 @@ + + WITH + +analytics_base AS ( + SELECT * + FROM `test-project.test-dataset.test-table` + WHERE repoid = @repoid + AND branch_name = @branch + AND timestamp BETWEEN + (CURRENT_TIMESTAMP() - INTERVAL @interval_start DAY) AND + (CURRENT_TIMESTAMP() - INTERVAL @interval_end DAY) +) + + +SELECT + name, + classname, + testsuite, + ANY_VALUE(computed_name) AS computed_name, + COUNT(DISTINCT IF(outcome = 1 OR outcome = 3, commit_sha, NULL)) AS cwf, + AVG(duration_seconds) AS avg_duration, + MAX_BY(duration_seconds, timestamp) AS last_duration, + SUM(IF(outcome = 0, 1, 0)) AS pass_count, + SUM(IF(outcome = 1, 1, 0)) AS fail_count, + SUM(IF(outcome = 2, 1, 0)) AS skip_count, + SUM(IF(outcome = 3, 1, 0)) AS flaky_fail_count, + MAX(timestamp) AS updated_at, + ARRAY_AGG(DISTINCT unnested_flags) AS flags +FROM analytics_base, UNNEST(flags) AS unnested_flags +GROUP BY name, classname, testsuite + + \ No newline at end of file diff --git a/ta_storage/tests/snapshots/bq__cache_analytics__0.json b/ta_storage/tests/snapshots/bq__cache_analytics__0.json new file mode 100644 index 000000000..a31e08337 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__cache_analytics__0.json @@ -0,0 +1,40 @@ +{ + "name": [ + "test_something" + ], + "classname": [ + "TestClass" + ], + "testsuite": [ + "test_suite" + ], + "computed_name": [ + "TestClass.test_something" + ], + "flags": [ + [ + "unit" + ] + ], + "avg_duration": [ + 1.5 + ], + "fail_count": [ + 2 + ], + "flaky_fail_count": [ + 1 + ], + "pass_count": [ + 10 + ], + "skip_count": [ + 1 + ], + "commits_where_fail": [ + 2 + ], + "last_duration": [ + 1.2 + ] +} diff --git a/ta_storage/tests/snapshots/bq__cache_analytics__1.txt b/ta_storage/tests/snapshots/bq__cache_analytics__1.txt new file mode 100644 index 000000000..6785389e7 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__cache_analytics__1.txt @@ -0,0 +1,32 @@ + + WITH + +analytics_base AS ( + SELECT * + FROM `test-project.test-dataset.test-table` + WHERE repoid = @repoid + AND branch_name = @branch + AND timestamp BETWEEN + (CURRENT_TIMESTAMP() - INTERVAL @interval_start DAY) AND + (CURRENT_TIMESTAMP() - INTERVAL @interval_end DAY) +) + + +SELECT + name, + classname, + testsuite, + ANY_VALUE(computed_name) AS computed_name, + COUNT(DISTINCT IF(outcome = 1 OR outcome = 3, commit_sha, NULL)) AS cwf, + AVG(duration_seconds) AS avg_duration, + MAX_BY(duration_seconds, timestamp) AS last_duration, + SUM(IF(outcome = 0, 1, 0)) AS pass_count, + SUM(IF(outcome = 1, 1, 0)) AS fail_count, + SUM(IF(outcome = 2, 1, 0)) AS skip_count, + SUM(IF(outcome = 3, 1, 0)) AS flaky_fail_count, + MAX(timestamp) AS updated_at, + ARRAY_AGG(DISTINCT unnested_flags) AS flags +FROM analytics_base, UNNEST(flags) AS unnested_flags +GROUP BY name, classname, testsuite + + \ No newline at end of file diff --git a/ta_storage/tests/snapshots/bq__pr_comment_agg__0.txt b/ta_storage/tests/snapshots/bq__pr_comment_agg__0.txt new file mode 100644 index 000000000..5c83f5064 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__pr_comment_agg__0.txt @@ -0,0 +1,51 @@ + + WITH + +ranked_data AS ( + SELECT + *, + ROW_NUMBER() OVER ( + PARTITION BY + name, + classname, + testsuite, + flags_hash + ORDER BY timestamp DESC + ) AS row_num + FROM + `test-project.test-dataset.test-table` + WHERE + repoid = @repoid + AND commit_sha = @commit_sha +) +, + +latest_instances AS ( + SELECT + * + FROM + ranked_data + WHERE + row_num = 1 +) + + +SELECT + * +FROM ( + SELECT + commit_sha, + outcome + FROM + latest_instances +) PIVOT ( + COUNT(*) AS ct + FOR outcome IN ( + 0 as passed, + 1 as failed, + 2 as skipped, + 3 as flaky_failed + ) +) + + \ No newline at end of file diff --git a/ta_storage/tests/snapshots/bq__pr_comment_agg__1.json b/ta_storage/tests/snapshots/bq__pr_comment_agg__1.json new file mode 100644 index 000000000..4a79117dc --- /dev/null +++ b/ta_storage/tests/snapshots/bq__pr_comment_agg__1.json @@ -0,0 +1,7 @@ +{ + "commit_sha": "abc123", + "passed_ct": 10, + "failed_ct": 2, + "skipped_ct": 1, + "flaky_failed_ct": 1 +} diff --git a/ta_storage/tests/snapshots/bq__pr_comment_fail__0.txt b/ta_storage/tests/snapshots/bq__pr_comment_fail__0.txt new file mode 100644 index 000000000..d0dc76f00 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__pr_comment_fail__0.txt @@ -0,0 +1,45 @@ + + WITH + +ranked_data AS ( + SELECT + *, + ROW_NUMBER() OVER ( + PARTITION BY + name, + classname, + testsuite, + flags_hash + ORDER BY timestamp DESC + ) AS row_num + FROM + `test-project.test-dataset.test-table` + WHERE + repoid = @repoid + AND commit_sha = @commit_sha +) +, + +latest_instances AS ( + SELECT + * + FROM + ranked_data + WHERE + row_num = 1 +) + + +SELECT + computed_name, + failure_message, + test_id, + flags_hash, + duration_seconds, + upload_id +FROM + latest_instances +WHERE + outcome = 1 + + \ No newline at end of file diff --git a/ta_storage/tests/snapshots/bq__pr_comment_fail__1.json b/ta_storage/tests/snapshots/bq__pr_comment_fail__1.json new file mode 100644 index 000000000..40410909d --- /dev/null +++ b/ta_storage/tests/snapshots/bq__pr_comment_fail__1.json @@ -0,0 +1,12 @@ +[ + { + "computed_name": "TestClass.test_something", + "failure_message": "assertion failed", + "id": [ + "746573745f6964", + "666c6167735f68617368" + ], + "duration_seconds": 1.5, + "upload_id": 1 + } +] diff --git a/ta_storage/tests/snapshots/bq__write_flakes__0.txt b/ta_storage/tests/snapshots/bq__write_flakes__0.txt new file mode 100644 index 000000000..8f62c2543 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__write_flakes__0.txt @@ -0,0 +1,24 @@ + + +SELECT + timestamp, + test_id, + outcome, + branch_name, + flags_hash, +FROM + `test-project.test-dataset.test-table` +WHERE + upload_id = @upload_id + AND ( + outcome = 1 + OR outcome = 3 + OR EXISTS ( + SELECT 1 + FROM UNNEST(@flake_ids) AS flake_id + WHERE flake_id.candidate_test_id = test_id + AND IF(flake_id.candidate_flags_hash IS NULL, flags_hash IS NULL, flake_id.candidate_flags_hash = flags_hash) + ) + ) + + \ No newline at end of file diff --git a/ta_storage/tests/snapshots/bq__write_flakes__1.json b/ta_storage/tests/snapshots/bq__write_flakes__1.json new file mode 100644 index 000000000..f9431ea60 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__write_flakes__1.json @@ -0,0 +1,12 @@ +[ + { + "repoid": 1, + "test_id": "746573745f6964", + "fail_count": 1, + "count": 1, + "recent_passes_count": 0, + "start_date": "2025-01-01T00:00:00+00:00", + "end_date": null, + "flags_id": null + } +] diff --git a/ta_storage/tests/snapshots/bq__write_testruns__0.json b/ta_storage/tests/snapshots/bq__write_testruns__0.json new file mode 100644 index 000000000..c22ec0947 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__write_testruns__0.json @@ -0,0 +1,22 @@ +[ + { + "timestamp": "1735689600000000", + "name": "test_something", + "classname": "TestClass", + "testsuite": "test_suite", + "computed_name": "TestClass.test_something", + "outcome": "PASSED", + "duration_seconds": 1.5, + "repoid": "1", + "commit_sha": "abc123", + "branch_name": "main", + "flags": [ + "unit" + ], + "filename": "test_file.py", + "framework": "pytest", + "upload_id": "1", + "flags_hash": "XLlcxAYPeBI=", + "test_id": "qbBVI/QFDPj9mlIXcoeGQA==" + } +] diff --git a/ta_storage/tests/snapshots/bq__write_testruns_with_flake__0.json b/ta_storage/tests/snapshots/bq__write_testruns_with_flake__0.json new file mode 100644 index 000000000..36254aeb0 --- /dev/null +++ b/ta_storage/tests/snapshots/bq__write_testruns_with_flake__0.json @@ -0,0 +1,23 @@ +[ + { + "timestamp": "1735689600000000", + "name": "test_something", + "classname": "TestClass", + "testsuite": "test_suite", + "computed_name": "TestClass.test_something", + "outcome": "FAILED", + "failure_message": "assertion failed", + "duration_seconds": 1.5, + "repoid": "1", + "commit_sha": "abc123", + "branch_name": "main", + "flags": [ + "unit" + ], + "filename": "test_file.py", + "framework": "pytest", + "upload_id": "1", + "flags_hash": "XLlcxAYPeBI=", + "test_id": "qbBVI/QFDPj9mlIXcoeGQA==" + } +] diff --git a/ta_storage/tests/test_bq.py b/ta_storage/tests/test_bq.py index 208190098..98daa49f7 100644 --- a/ta_storage/tests/test_bq.py +++ b/ta_storage/tests/test_bq.py @@ -1,19 +1,27 @@ -from __future__ import annotations - -from datetime import datetime, timedelta, timezone -from unittest.mock import MagicMock, patch +from datetime import datetime, timezone +from unittest.mock import ANY, MagicMock, patch +import polars as pl import pytest import test_results_parser -import time_machine +from google.cloud.bigquery import ( + ArrayQueryParameter, + ScalarQueryParameter, + StructQueryParameterType, +) +from google.protobuf.json_format import MessageToDict +from shared.django_apps.reports.tests.factories import ( + UploadFactory, +) +from shared.django_apps.test_analytics.models import Flake +from time_machine import travel import generated_proto.testrun.ta_testrun_pb2 as ta_testrun_pb2 -from database.tests.factories import RepositoryFlagFactory, UploadFactory from ta_storage.bq import BQDriver from ta_storage.utils import calc_flags_hash, calc_test_id -@pytest.fixture +@pytest.fixture() def mock_bigquery_service(): with patch("ta_storage.bq.get_bigquery_service") as mock: service = MagicMock() @@ -21,465 +29,353 @@ def mock_bigquery_service(): yield service -def test_bigquery_driver(dbsession, mock_bigquery_service): - bq = BQDriver() +@pytest.fixture() +def mock_config(mock_configuration): + mock_configuration._params["services"]["gcp"] = { + "project_id": "test-project", + } + mock_configuration._params["services"]["bigquery"] = { + "dataset_name": "test-dataset", + "testrun_table_name": "test-table", + } - upload = UploadFactory() - dbsession.add(upload) - dbsession.flush() - repo_flag_1 = RepositoryFlagFactory( - repository=upload.report.commit.repository, flag_name="flag1" - ) - repo_flag_2 = RepositoryFlagFactory( - repository=upload.report.commit.repository, flag_name="flag2" - ) - dbsession.add(repo_flag_1) - dbsession.add(repo_flag_2) - dbsession.flush() +def read_table(mock_storage, bucket: str, storage_path: str): + decompressed_table: bytes = mock_storage.read_file(bucket, storage_path) + return pl.read_ipc(decompressed_table) + - upload.flags.append(repo_flag_1) - upload.flags.append(repo_flag_2) - dbsession.flush() +@pytest.mark.django_db(transaction=True, databases=["test_analytics"]) +def test_write_testruns(mock_bigquery_service, mock_config, snapshot): + driver = BQDriver(repo_id=1) + timestamp = int( + datetime.fromisoformat("2025-01-01T00:00:00Z").timestamp() * 1000000 + ) - test_data: list[test_results_parser.Testrun] = [ + testruns: list[test_results_parser.Testrun] = [ { - "name": "test_name", - "classname": "test_class", + "name": "test_something", + "classname": "TestClass", "testsuite": "test_suite", - "duration": 100.0, + "computed_name": "TestClass.test_something", "outcome": "pass", - "build_url": "https://example.com/build/123", - "filename": "test_file", - "computed_name": "test_computed_name", "failure_message": None, - }, - { - "name": "test_name2", - "classname": "test_class2", - "testsuite": "test_suite2", - "duration": 100.0, - "outcome": "failure", - "build_url": "https://example.com/build/123", - "filename": "test_file2", - "computed_name": "test_computed_name2", - "failure_message": "test_failure_message", - }, + "duration": 1.5, + "filename": "test_file.py", + "build_url": None, + } ] - timestamp = int(datetime.now().timestamp() * 1000000) - - bq.write_testruns( - timestamp, - upload.report.commit.repoid, - upload.report.commit.commitid, - upload.report.commit.branch, - upload, - "pytest", - test_data, + driver.write_testruns( + timestamp=timestamp, + commit_sha="abc123", + branch_name="main", + upload_id=1, + flag_names=["unit"], + framework="pytest", + testruns=testruns, ) - flags_hash = calc_flags_hash(upload.flag_names) - - # Verify the BigQuery service was called correctly mock_bigquery_service.write.assert_called_once_with( - bq.dataset_name, - bq.testrun_table_name, - ta_testrun_pb2, - [ - ta_testrun_pb2.TestRun( - timestamp=timestamp, - name="test_name", - classname="test_class", - testsuite="test_suite", - duration_seconds=100.0, - outcome=ta_testrun_pb2.TestRun.Outcome.PASSED, - filename="test_file", - computed_name="test_computed_name", - failure_message=None, - repoid=upload.report.commit.repoid, - commit_sha=upload.report.commit.commitid, - framework="pytest", - branch_name=upload.report.commit.branch, - flags=["flag1", "flag2"], - upload_id=upload.id_, - flags_hash=flags_hash, - test_id=calc_test_id("test_name", "test_class", "test_suite"), - ).SerializeToString(), - ta_testrun_pb2.TestRun( - timestamp=timestamp, - name="test_name2", - classname="test_class2", - testsuite="test_suite2", - duration_seconds=100.0, - outcome=ta_testrun_pb2.TestRun.Outcome.FAILED, - filename="test_file2", - computed_name="test_computed_name2", - failure_message="test_failure_message", - repoid=upload.report.commit.repoid, - commit_sha=upload.report.commit.commitid, - framework="pytest", - branch_name=upload.report.commit.branch, - flags=["flag1", "flag2"], - upload_id=upload.id_, - flags_hash=flags_hash, - test_id=calc_test_id("test_name2", "test_class2", "test_suite2"), - ).SerializeToString(), - ], + "test-dataset", "test-table", ta_testrun_pb2, ANY ) - -def populate_pr_comment_testruns(bq: BQDriver): - testruns = [] - - for i in range(3): - upload = UploadFactory() - upload.report.commit.commitid = "abcde" - upload.report.commit.branch = "feature_branch" - upload.report.commit.repoid = 2 - upload.flags.append(RepositoryFlagFactory(flag_name=f"flag_{i}")) - - for j in range(3): - name = f"test_{j}" - classname = f"class_{j}" - testsuite = "suite_feature" - - testrun: test_results_parser.Testrun = { - "name": name, - "classname": classname, - "testsuite": testsuite, - "duration": float(j % 5), - "outcome": "pass" if j % 2 == 0 else "failure", - "filename": None, - "computed_name": f"pr_computed_name_{j}", - "failure_message": None if j % 2 == 0 else "hi", - "build_url": None, - } - - testruns.append(testrun) - - bq.write_testruns( - None, 2, "abcde", "feature_branch", upload, "pytest", testruns + testruns_written = [ + MessageToDict( + ta_testrun_pb2.TestRun.FromString(testrun_bytes), + preserving_proto_field_name=True, ) + for testrun_bytes in mock_bigquery_service.mock_calls[0][1][3] + ] + assert snapshot("json") == sorted(testruns_written, key=lambda x: x["name"]) -@pytest.mark.skip(reason="need creds") -def test_bq_pr_comment(): - bq = BQDriver() +@pytest.mark.django_db(transaction=True, databases=["test_analytics"]) +def test_write_testruns_with_flake(mock_bigquery_service, mock_config, snapshot): + driver = BQDriver(repo_id=1) + timestamp = int( + datetime.fromisoformat("2025-01-01T00:00:00Z").timestamp() * 1000000 + ) - if ( - bq.bq_service.query( - "select * from `test_dataset.testruns` where repoid = 2 limit 1" - ) - == [] - ): - populate_pr_comment_testruns(bq) + flake = Flake.objects.create( + repoid=1, + test_id=calc_test_id("test_suite", "TestClass", "test_something"), + flags_id=calc_flags_hash(["unit"]), + fail_count=1, + count=1, + recent_passes_count=1, + start_date=datetime.now(timezone.utc), + ) + flake.save() - pr_agg = bq.pr_comment_agg(repoid=2, commit_sha="abcde") - assert pr_agg == [ + testruns: list[test_results_parser.Testrun] = [ { - "commit_sha": "abcde", - "ct_passed": 6, - "ct_failed": 3, - "ct_skipped": 0, - "ct_flaky_failed": 0, + "name": "test_something", + "classname": "TestClass", + "testsuite": "test_suite", + "computed_name": "TestClass.test_something", + "outcome": "failure", + "failure_message": "assertion failed", + "duration": 1.5, + "filename": "test_file.py", + "build_url": None, } ] - pr_fail = bq.pr_comment_fail(repoid=2, commit_sha="abcde") - assert len(pr_fail) == 3 - assert {t["computed_name"] for t in pr_fail} == { - "pr_computed_name_1", - } - assert {t["failure_message"] for t in pr_fail} == {"hi"} - assert {tuple(t["flags"]) for t in pr_fail} == { - ("flag_1",), - ("flag_2",), - ("flag_0",), - } + driver.write_testruns( + timestamp=timestamp, + commit_sha="abc123", + branch_name="main", + upload_id=1, + flag_names=["unit"], + framework="pytest", + testruns=testruns, + ) + mock_bigquery_service.write.assert_called_once_with( + "test-dataset", "test-table", ta_testrun_pb2, ANY + ) -def populate_testruns_for_upload_testruns(dbsession, bq: BQDriver): - testruns = [] + testruns_written = [ + MessageToDict( + ta_testrun_pb2.TestRun.FromString(testrun_bytes), + preserving_proto_field_name=True, + ) + for testrun_bytes in mock_bigquery_service.mock_calls[0][1][3] + ] + assert snapshot("json") == sorted(testruns_written, key=lambda x: x["name"]) - upload = UploadFactory() - upload.id_ = 1 - dbsession.add(upload) - dbsession.flush() - testruns: list[test_results_parser.Testrun] = [ - { # this test is flaky failure - "name": "test_0", - "classname": "class_0", - "testsuite": "suite_upload", - "duration": 0.0, - "outcome": "failure", - "filename": None, - "computed_name": "upload_computed_name_0", - "failure_message": None, - "build_url": None, - }, - { # this test is just a failure - "name": "test_1", - "classname": "class_1", - "testsuite": "suite_upload", - "duration": 0.0, - "outcome": "failure", - "filename": None, - "computed_name": "upload_computed_name_1", - "failure_message": None, - "build_url": None, - }, - { # this test is a pass but also flaky - "name": "test_2", - "classname": "class_2", - "testsuite": "suite_upload", - "duration": 0.0, - "outcome": "pass", - "filename": None, - "computed_name": "upload_computed_name_2", - "failure_message": None, - "build_url": None, - }, - { # this test should be ignored - "name": "test_3", - "classname": "class_3", - "testsuite": "suite_upload", - "duration": 0.0, - "outcome": "pass", - "filename": None, - "computed_name": "upload_computed_name_3", - "failure_message": None, - "build_url": None, - }, +def test_pr_comment_agg(mock_bigquery_service, mock_config, snapshot): + driver = BQDriver(repo_id=1) + mock_bigquery_service.query.return_value = [ + { + "commit_sha": "abc123", + "passed": 10, + "failed": 2, + "skipped": 1, + "flaky_failed": 1, + } ] - bq.write_testruns(None, 3, "abcde", "feature_branch", upload, "pytest", testruns) + result = driver.pr_comment_agg("abc123") + + mock_bigquery_service.query.assert_called_once() + query, params = mock_bigquery_service.query.call_args[0] + assert snapshot("txt") == query + assert params == [ + ScalarQueryParameter("repoid", "INT64", 1), + ScalarQueryParameter("commit_sha", "STRING", "abc123"), + ] + assert snapshot("json") == result -@pytest.mark.skip(reason="need creds") -def test_bq_testruns_for_upload(dbsession): - bq = BQDriver( +def test_pr_comment_fail(mock_bigquery_service, mock_config, snapshot): + driver = BQDriver(repo_id=1) + mock_bigquery_service.query.return_value = [ { - calc_test_id("test_0", "class_0", "suite_upload"), - calc_test_id("test_2", "class_2", "suite_upload"), + "computed_name": "TestClass.test_something", + "failure_message": "assertion failed", + "test_id": b"test_id", + "flags_hash": b"flags_hash", + "duration_seconds": 1.5, + "upload_id": 1, } - ) + ] - if ( - bq.bq_service.query( - "select * from `test_dataset.testruns` where repoid = 3 limit 1" - ) - == [] - ): - populate_testruns_for_upload_testruns(dbsession, bq) + result = driver.pr_comment_fail("abc123") - testruns_for_upload = bq.testruns_for_upload( - upload_id=1, - test_ids=[ - calc_test_id("test_0", "class_0", "suite_upload"), - calc_test_id("test_2", "class_2", "suite_upload"), - ], - ) - - assert {t["test_id"] for t in testruns_for_upload} == { - calc_test_id("test_0", "class_0", "suite_upload"), - calc_test_id("test_2", "class_2", "suite_upload"), - calc_test_id("test_1", "class_1", "suite_upload"), - } - - assert {t["outcome"] for t in testruns_for_upload} == {3, 1, 0} + # Verify the query parameters + mock_bigquery_service.query.assert_called_once() + query, params = mock_bigquery_service.query.call_args[0] + assert snapshot("txt") == query + assert params == [ + ScalarQueryParameter("repoid", "INT64", 1), + ScalarQueryParameter("commit_sha", "STRING", "abc123"), + ] + result[0]["id"] = [result[0]["id"][0].hex(), result[0]["id"][1].hex()] + assert snapshot("json") == result -def populate_analytics_testruns(bq: BQDriver): - upload_0 = UploadFactory() - upload_0.report.commit.commitid = "abcde" - upload_0.report.commit.branch = "feature_branch" - upload_0.report.commit.repoid = 1 - upload_0.flags.append(RepositoryFlagFactory(flag_name="flag_0")) +@travel("2025-01-01T00:00:00Z", tick=False) +@pytest.mark.django_db(transaction=True, databases=["default", "test_analytics"]) +def test_write_flakes(mock_bigquery_service, mock_config, snapshot): + driver = BQDriver(repo_id=1) - upload_1 = UploadFactory() - upload_1.report.commit.commitid = "abcde" - upload_1.report.commit.branch = "feature_branch" - upload_1.report.commit.repoid = 1 - upload_1.flags.append(RepositoryFlagFactory(flag_name="flag_1")) + upload = UploadFactory.create() + upload.save() - testruns: list[test_results_parser.Testrun] = [ + mock_bigquery_service.query.return_value = [ { - "name": "interval_start", - "classname": "class_0", - "testsuite": "suite_upload", - "duration": 20000.0, - "outcome": "failure", - "filename": None, - "computed_name": "upload_computed_name_0", - "failure_message": None, - "build_url": None, - }, + "branch_name": "main", + "timestamp": int(datetime.now().timestamp() * 1000000), + "outcome": ta_testrun_pb2.TestRun.Outcome.FAILED, + "test_id": b"test_id", + "flags_hash": b"flags_hash", + } ] - timestamp = int((datetime.now() - timedelta(days=50)).timestamp() * 1000000) - - bq.write_testruns( - timestamp, 1, "interval_start", "feature_branch", upload_0, "pytest", testruns - ) + driver.write_flakes([upload]) + + mock_bigquery_service.query.assert_called_once() + query, params = mock_bigquery_service.query.call_args[0] + assert snapshot("txt") == query + assert params == [ + ScalarQueryParameter("upload_id", "INT64", upload.id), + ArrayQueryParameter( + "flake_ids", + StructQueryParameterType( + ScalarQueryParameter("test_id", "STRING", "test_id"), + ScalarQueryParameter("flags_id", "STRING", "flags_id"), + ), + [], + ), + ] - testruns: list[test_results_parser.Testrun] = [ + flakes = Flake.objects.all() + flake_data = [ { - "name": "interval_end", - "classname": "class_0", - "testsuite": "suite_upload", - "duration": 20000.0, - "outcome": "failure", - "filename": None, - "computed_name": "upload_computed_name_0", - "failure_message": None, - "build_url": None, - }, + "repoid": flake.repoid, + "test_id": flake.test_id.hex(), + "fail_count": flake.fail_count, + "count": flake.count, + "recent_passes_count": flake.recent_passes_count, + "start_date": flake.start_date.isoformat() if flake.start_date else None, + "end_date": flake.end_date.isoformat() if flake.end_date else None, + "flags_id": flake.flags_id.hex() if flake.flags_id else None, + } + for flake in flakes ] + assert snapshot("json") == sorted(flake_data, key=lambda x: x["test_id"]) - timestamp = int((datetime.now() - timedelta(days=1)).timestamp() * 1000000) - bq.write_testruns( - timestamp, 1, "interval_end", "feature_branch", upload_0, "pytest", testruns - ) +def test_analytics(mock_bigquery_service, mock_config, snapshot): + driver = BQDriver(repo_id=1) - testruns: list[test_results_parser.Testrun] = [ - { - "name": "test_0", - "classname": "class_0", - "testsuite": "suite_upload", - "duration": 10.0, - "outcome": "failure", - "filename": None, - "computed_name": "upload_computed_name_0", - "failure_message": None, - "build_url": None, - }, - { - "name": "test_1", - "classname": "class_1", - "testsuite": "suite_upload", - "duration": 10.0, - "outcome": "pass", - "filename": None, - "computed_name": "upload_computed_name_1", - "failure_message": None, - "build_url": None, - }, + _ = driver.analytics(1, 30, 0) + query, params = mock_bigquery_service.query.call_args[0] + assert snapshot("txt") == query + assert params == [ + ScalarQueryParameter("repoid", "INT64", 1), + ScalarQueryParameter("interval_start", "INT64", 30), + ScalarQueryParameter("interval_end", "INT64", 0), ] - timestamp = int((datetime.now() - timedelta(days=20)).timestamp() * 1000000) - bq.write_testruns( - timestamp, 1, "commit_1", "feature_branch", upload_0, "pytest", testruns - ) +def test_analytics_with_branch(mock_bigquery_service, mock_config, snapshot): + driver = BQDriver(repo_id=1) - testruns: list[test_results_parser.Testrun] = [ - { - "name": "test_1", - "classname": "class_1", - "testsuite": "suite_upload", - "duration": 10.0, - "outcome": "failure", - "filename": None, - "computed_name": "upload_computed_name_1", - "failure_message": None, - "build_url": None, - }, + _ = driver.analytics(1, 30, 0, "main") + query, params = mock_bigquery_service.query.call_args[0] + assert snapshot("txt") == query + assert params == [ + ScalarQueryParameter("repoid", "INT64", 1), + ScalarQueryParameter("interval_start", "INT64", 30), + ScalarQueryParameter("interval_end", "INT64", 0), + ScalarQueryParameter("branch", "STRING", "main"), ] - timestamp = int((datetime.now() - timedelta(days=20)).timestamp() * 1000000) - bq.write_testruns( - timestamp, 1, "commit_1", "feature_branch", upload_1, "pytest", testruns - ) +@travel("2025-01-01T00:00:00Z", tick=False) +def test_cache_analytics(mock_bigquery_service, mock_config, mock_storage, snapshot): + driver = BQDriver(repo_id=1) - bq = BQDriver( + # Mock the analytics query result with datetime + mock_bigquery_service.query.return_value = [ { - calc_test_id("test_1", "class_1", "suite_upload"), + "name": "test_something", + "classname": "TestClass", + "testsuite": "test_suite", + "computed_name": "TestClass.test_something", + "flags": ["unit"], + "avg_duration": 1.5, + "fail_count": 2, + "flaky_fail_count": 1, + "pass_count": 10, + "skip_count": 1, + "commits_where_fail": 2, + "last_duration": 1.2, + "updated_at": datetime.fromisoformat("2025-01-01T00:00:00+00:00"), } - ) - - testruns: list[test_results_parser.Testrun] = [ - { - "name": "test_0", - "classname": "class_0", - "testsuite": "suite_upload", - "duration": 20.0, - "outcome": "pass", - "filename": None, - "computed_name": "upload_computed_name_0", - "failure_message": None, - "build_url": None, - }, - { - "name": "test_1", - "classname": "class_1", - "testsuite": "suite_upload", - "duration": 10.0, - "outcome": "failure", - "filename": None, - "computed_name": "upload_computed_name_1", - "failure_message": None, - "build_url": None, - }, ] - timestamp = int((datetime.now() - timedelta(days=10)).timestamp() * 1000000) + buckets = ["bucket1", "bucket2"] + branch = "main" + driver.cache_analytics(buckets, branch) - bq.write_testruns( - timestamp, 1, "commit_2", "feature_branch", upload_1, "pytest", testruns - ) + assert mock_bigquery_service.query.call_count == 6 + expected_intervals = [ + (1, None), + (2, 1), + (7, None), + (14, 7), + (30, None), + (60, 30), + ] -@pytest.mark.skip(reason="need creds") -@time_machine.travel(datetime.now(tz=timezone.utc), tick=False) -def test_bq_analytics(): - bq = BQDriver() - - if ( - bq.bq_service.query( - "select * from `test_dataset.testruns` where repoid = 1 limit 1" - ) - == [] - ): - populate_analytics_testruns(bq) + expected_dict = { + "name": [], + "classname": [], + "testsuite": [], + "computed_name": [], + "flags": [], + "avg_duration": [], + "fail_count": [], + "flaky_fail_count": [], + "pass_count": [], + "skip_count": [], + "commits_where_fail": [], + "last_duration": [], + } - testruns_for_upload = bq.analytics(1, 30, 7, "feature_branch") + table_dicts = [] + for bucket in buckets: + for interval_start, interval_end in expected_intervals: + storage_key = ( + f"ta_rollups/{driver.repo_id}/{branch}/{interval_start}" + if interval_end is None + else f"ta_rollups/{driver.repo_id}/{branch}/{interval_start}_{interval_end}" + ) + table = read_table(mock_storage, bucket, storage_key) + table_dict = table.to_dict(as_series=False) + table_dicts.append(table_dict) + + assert snapshot("json") == { + "name": table_dicts[0]["name"], + "classname": table_dicts[0]["classname"], + "testsuite": table_dicts[0]["testsuite"], + "computed_name": table_dicts[0]["computed_name"], + "flags": table_dicts[0]["flags"], + "avg_duration": table_dicts[0]["avg_duration"], + "fail_count": table_dicts[0]["fail_count"], + "flaky_fail_count": table_dicts[0]["flaky_fail_count"], + "pass_count": table_dicts[0]["pass_count"], + "skip_count": table_dicts[0]["skip_count"], + "commits_where_fail": table_dicts[0]["commits_where_fail"], + "last_duration": table_dicts[0]["last_duration"], + } - assert sorted( - [(x | {"flags": sorted(x["flags"])}) for x in testruns_for_upload], - key=lambda x: x["name"], - ) == [ - { - "name": "test_0", - "classname": "class_0", - "testsuite": "suite_upload", - "computed_name": "upload_computed_name_0", - "cwf": 1, - "avg_duration": 15.0, - "last_duration": 20.0, - "pass_count": 1, - "fail_count": 1, - "skip_count": 0, - "flaky_fail_count": 0, - "updated_at": datetime.now(tz=timezone.utc) - timedelta(days=10), - "flags": ["flag_0", "flag_1"], - }, - { - "name": "test_1", - "classname": "class_1", - "testsuite": "suite_upload", - "computed_name": "upload_computed_name_1", - "cwf": 2, - "avg_duration": 10.0, - "last_duration": 10.0, - "pass_count": 1, - "fail_count": 1, - "skip_count": 0, - "flaky_fail_count": 1, - "updated_at": datetime.now(tz=timezone.utc) - timedelta(days=10), - "flags": ["flag_0", "flag_1"], - }, - ] + first_dict = table_dicts[0] + for table_dict in table_dicts[1:]: + assert table_dict == first_dict + + queries = [] + params = [] + for args in mock_bigquery_service.query.call_args_list: + queries.append(args[0][0]) + params.append(args[0][1]) + + first_query = queries[0] + for query in queries[1:]: + assert query == first_query + + assert snapshot("txt") == first_query + + for i, (interval_start, interval_end) in enumerate(expected_intervals): + assert params[i] == [ + ScalarQueryParameter("repoid", "INT64", 1), + ScalarQueryParameter("interval_start", "INT64", interval_start), + ScalarQueryParameter("interval_end", "INT64", interval_end), + ScalarQueryParameter("branch", "STRING", branch), + ] diff --git a/ta_storage/tests/test_pg.py b/ta_storage/tests/test_pg.py index 370cd1b42..a6af0654d 100644 --- a/ta_storage/tests/test_pg.py +++ b/ta_storage/tests/test_pg.py @@ -1,11 +1,13 @@ from database.models import DailyTestRollup, Test, TestFlagBridge, TestInstance -from database.tests.factories import RepositoryFlagFactory, UploadFactory +from database.tests.factories import ( + RepositoryFlagFactory, + UploadFactory, +) +from database.tests.factories.reports import TestFactory, TestInstanceFactory from ta_storage.pg import PGDriver def test_pg_driver(dbsession): - pg = PGDriver(dbsession, set()) - upload = UploadFactory() dbsession.add(upload) dbsession.flush() @@ -24,12 +26,13 @@ def test_pg_driver(dbsession): upload.flags.append(repo_flag_2) dbsession.flush() + pg = PGDriver(upload.report.commit.repoid, dbsession, None) pg.write_testruns( None, - upload.report.commit.repoid, upload.report.commit.id, upload.report.commit.branch, - upload, + upload.id_, + upload.flag_names, "pytest", [ { @@ -61,3 +64,105 @@ def test_pg_driver(dbsession): assert dbsession.query(TestInstance).count() == 2 assert dbsession.query(TestFlagBridge).count() == 4 assert dbsession.query(DailyTestRollup).count() == 2 + + +def test_pg_driver_pr_comment_agg(dbsession): + # Create test data + upload = UploadFactory() + dbsession.add(upload) + dbsession.flush() + + test = TestFactory( + repoid=upload.report.commit.repoid, + name="test_name", + testsuite="test_suite", + computed_name="test_computed_name", + ) + dbsession.add(test) + dbsession.flush() + + test_instance = TestInstanceFactory( + test=test, + upload=upload, + outcome="pass", + commitid=upload.report.commit.commitid, + branch=upload.report.commit.branch, + repoid=upload.report.commit.repoid, + ) + dbsession.add(test_instance) + dbsession.flush() + + test_instance_2 = TestInstanceFactory( + test=test, + upload=upload, + outcome="failure", + commitid=upload.report.commit.commitid, + branch=upload.report.commit.branch, + repoid=upload.report.commit.repoid, + ) + dbsession.add(test_instance_2) + dbsession.flush() + + test_instance_3 = TestInstanceFactory( + test=test, + upload=upload, + outcome="skip", + commitid=upload.report.commit.commitid, + branch=upload.report.commit.branch, + repoid=upload.report.commit.repoid, + ) + dbsession.add(test_instance_3) + dbsession.flush() + + pg = PGDriver(upload.report.commit.repoid, dbsession, None) + result = pg.pr_comment_agg(upload.report.commit.commitid) + + assert result == { + "commit_sha": upload.report.commit.commitid, + "passed_ct": 1, + "failed_ct": 1, + "skipped_ct": 1, + "flaky_failed_ct": 0, + } + + +def test_pg_driver_pr_comment_fail(dbsession): + # Create test data + upload = UploadFactory() + dbsession.add(upload) + upload.id_ = 3 + dbsession.flush() + + test = TestFactory( + repoid=upload.report.commit.repoid, + name="test_name", + testsuite="test_suite", + computed_name="test_computed_name", + ) + dbsession.add(test) + dbsession.flush() + + test_instance = TestInstanceFactory( + test=test, + upload=upload, + outcome="failure", + failure_message="Test failed with error", + commitid=upload.report.commit.commitid, + branch=upload.report.commit.branch, + repoid=upload.report.commit.repoid, + ) + dbsession.add(test_instance) + dbsession.flush() + + pg = PGDriver(upload.report.commit.repoid, dbsession, None) + result = pg.pr_comment_fail(upload.report.commit.commitid) + + assert result == [ + { + "computed_name": "test_computed_name", + "duration_seconds": 1.5, + "failure_message": "Test failed with error", + "id": "id_1", + "upload_id": 3, + } + ] diff --git a/ta_storage/tests/test_pg_django.py b/ta_storage/tests/test_pg_django.py new file mode 100644 index 000000000..c73da8801 --- /dev/null +++ b/ta_storage/tests/test_pg_django.py @@ -0,0 +1,272 @@ +from datetime import datetime, timedelta + +import polars as pl +import pytest +from shared.django_apps.reports.models import DailyTestRollup, Flake +from shared.django_apps.reports.tests.factories import ( + DailyTestRollupFactory, + TestFactory, + TestInstanceFactory, + UploadFactory, +) + +from services.processing.flake_processing import FLAKE_EXPIRY_COUNT +from ta_storage.pg import PGDriver + + +def read_table(mock_storage, bucket: str, storage_path: str): + decompressed_table: bytes = mock_storage.read_file(bucket, storage_path) + return pl.read_ipc(decompressed_table) + + +@pytest.mark.django_db(transaction=True) +def test_pg_driver_cache_analytics(mock_storage): + # Create test data + upload = UploadFactory() + upload.save() + + test = TestFactory( + repository=upload.report.commit.repository, + name="test_name", + testsuite="test_suite", + computed_name="test_computed_name", + ) + test.save() + + test_instance = TestInstanceFactory( + test=test, + upload=upload, + outcome="pass", + duration_seconds=1.5, + commitid=upload.report.commit.commitid, + branch=upload.report.commit.branch, + repoid=upload.report.commit.repository.repoid, + ) + test_instance.save() + + # Create daily rollup data for different intervals + today = datetime.now().date() + + # Today's data + rollup_today = DailyTestRollupFactory( + test=test, + repoid=test.repository.repoid, + branch=upload.report.commit.branch, + date=today, + pass_count=10, + fail_count=2, + skip_count=1, + flaky_fail_count=1, + avg_duration_seconds=1.5, + last_duration_seconds=1.2, + latest_run=datetime.now(), + commits_where_fail=[upload.report.commit.commitid], + ) + rollup_today.save() + + # Yesterday's data + rollup_yesterday = DailyTestRollupFactory( + test=test, + repoid=test.repository.repoid, + branch=upload.report.commit.branch, + date=today - timedelta(days=1), + pass_count=5, + fail_count=3, + skip_count=0, + flaky_fail_count=2, + avg_duration_seconds=1.8, + last_duration_seconds=1.4, + latest_run=datetime.now() - timedelta(days=1), + commits_where_fail=[upload.report.commit.commitid], + ) + rollup_yesterday.save() + + # Data from 7 days ago + rollup_week = DailyTestRollupFactory( + test=test, + repoid=test.repository.repoid, + branch=upload.report.commit.branch, + date=today - timedelta(days=7), + pass_count=8, + fail_count=4, + skip_count=2, + flaky_fail_count=3, + avg_duration_seconds=1.6, + last_duration_seconds=1.3, + latest_run=datetime.now() - timedelta(days=7), + commits_where_fail=[upload.report.commit.commitid], + ) + rollup_week.save() + + # Data from 30 days ago + rollup_month = DailyTestRollupFactory( + test=test, + repoid=test.repository.repoid, + branch=upload.report.commit.branch, + date=today - timedelta(days=30), + pass_count=15, + fail_count=5, + skip_count=3, + flaky_fail_count=4, + avg_duration_seconds=1.7, + last_duration_seconds=1.5, + latest_run=datetime.now() - timedelta(days=30), + commits_where_fail=[upload.report.commit.commitid], + ) + rollup_month.save() + + buckets = ["bucket1", "bucket2"] + branch = upload.report.commit.branch + + pg = PGDriver(upload.report.commit.repository.repoid) + pg.cache_analytics(buckets, branch) + + rollups = DailyTestRollup.objects.filter( + repoid=test.repository.repoid, + branch=upload.report.commit.branch, + ) + + print(rollups) + + expected_intervals = [ + (1, None), # Today + (2, 1), # Yesterday + (7, None), # Last week + (14, 7), # Week before last + (30, None), # Last month + (60, 30), # Month before last + ] + + # Verify data for each interval in each bucket + for bucket in buckets: + for interval_start, interval_end in expected_intervals: + storage_key = ( + f"test_results/rollups/{upload.report.commit.repository.repoid}/{branch}/{interval_start}" + if interval_end is None + else f"test_results/rollups/{upload.report.commit.repository.repoid}/{branch}/{interval_start}_{interval_end}" + ) + table = read_table(mock_storage, bucket, storage_key) + table_dict = table.to_dict(as_series=False) + + print(table_dict) + # Verify data based on intervals + if (interval_start, interval_end) == (1, None): + # Today's data + assert table_dict["total_pass_count"] == [15] + assert table_dict["total_fail_count"] == [5] + elif (interval_start, interval_end) == (2, 1): + # Yesterday's data + assert table_dict["total_pass_count"] == [] + assert table_dict["total_fail_count"] == [] + elif (interval_start, interval_end) == (7, None): + # Last week's data (includes today) + assert table_dict["total_pass_count"] == [23] # 10 + 5 + 8 + assert table_dict["total_fail_count"] == [9] # 2 + 3 + 4 + elif (interval_start, interval_end) == (30, None): + # Last month's data (includes all) + assert table_dict["total_pass_count"] == [38] # 10 + 5 + 8 + 15 + assert table_dict["total_fail_count"] == [14] # 2 + 3 + 4 + 5 + elif (interval_start, interval_end) == (14, 7): + # Week before last (should be empty since we have no data in this range) + assert table_dict["total_pass_count"] == [] + assert table_dict["total_fail_count"] == [] + elif (interval_start, interval_end) == (60, 30): + # Month before last (should be empty) + assert table_dict["total_pass_count"] == [] + assert table_dict["total_fail_count"] == [] + + +@pytest.mark.django_db(transaction=True) +def test_pg_driver_write_flakes(mock_storage): + # Create test data + upload1 = UploadFactory() + upload1.save() + + test1 = TestFactory( + repository=upload1.report.commit.repository, + name="test_name1", + testsuite="test_suite", + computed_name="test_computed_name1", + ) + test1.save() + + test2 = TestFactory( + repository=upload1.report.commit.repository, + name="test_name2", + testsuite="test_suite", + computed_name="test_computed_name2", + ) + test2.save() + + # Create a pre-existing flake for test1 + flake1 = Flake.objects.create( + repository=test1.repository, + test=test1, + fail_count=1, + count=1, + recent_passes_count=0, + start_date=datetime.now(), + ) + flake1.save() + + # Create test instances that will update the existing flake + test_instance1 = TestInstanceFactory( + test=test1, + upload=upload1, + outcome="pass", # This should increment recent_passes_count + duration_seconds=1.5, + commitid=upload1.report.commit.commitid, + branch=upload1.report.commit.branch, + repoid=upload1.report.commit.repository.repoid, + ) + test_instance1.save() + + # Create test instances that will create a new flake + test_instance2 = TestInstanceFactory( + test=test2, + upload=upload1, + outcome="failure", # This should create a new flake + duration_seconds=1.8, + failure_message="Test failed", + commitid=upload1.report.commit.commitid, + branch=upload1.report.commit.branch, + repoid=upload1.report.commit.repository.repoid, + ) + test_instance2.save() + + # Create another upload with more test instances + upload2 = UploadFactory() + upload2.save() + + # Create test instances that will expire the flake + for _ in range(FLAKE_EXPIRY_COUNT - 1): # -1 because we already have one pass + test_instance = TestInstanceFactory( + test=test1, + upload=upload2, + outcome="pass", # These passes should eventually expire the flake + duration_seconds=1.5, + commitid=upload2.report.commit.commitid, + branch=upload2.report.commit.branch, + repoid=upload2.report.commit.repository.repoid, + ) + test_instance.save() + + pg = PGDriver(upload1.report.commit.repository.repoid) + pg.write_flakes([upload1, upload2]) + + # Verify the flakes + flakes = Flake.objects.filter(repository=test1.repository).order_by("test") + + # Verify the first flake (should be expired) + assert flakes[0].test == test1 + assert flakes[0].count == FLAKE_EXPIRY_COUNT + 1 # Original + all passes + assert flakes[0].fail_count == 1 # Unchanged + assert flakes[0].recent_passes_count == FLAKE_EXPIRY_COUNT + assert flakes[0].end_date is not None # Should be expired + + # Verify the second flake (newly created) + assert flakes[1].test == test2 + assert flakes[1].count == 1 + assert flakes[1].fail_count == 1 + assert flakes[1].recent_passes_count == 0 + assert flakes[1].end_date is None # Should not be expired diff --git a/ta_storage/utils.py b/ta_storage/utils.py index 0c3dd3e42..f15402fda 100644 --- a/ta_storage/utils.py +++ b/ta_storage/utils.py @@ -1,5 +1,9 @@ +import logging + import mmh3 +log = logging.getLogger(__name__) + def calc_test_id(name: str, classname: str, testsuite: str) -> bytes: h = mmh3.mmh3_x64_128() # assumes we're running on x64 machines @@ -13,6 +17,8 @@ def calc_test_id(name: str, classname: str, testsuite: str) -> bytes: def calc_flags_hash(flags: list[str]) -> bytes | None: flags_str = " ".join(sorted(flags)) # we know that flags cannot contain spaces + if not flags: + return None # returns a tuple of two int64 values # we only need the first one diff --git a/tasks/ta_processor.py b/tasks/ta_processor.py index e2b913bd3..64ed643b1 100644 --- a/tasks/ta_processor.py +++ b/tasks/ta_processor.py @@ -14,6 +14,7 @@ Upload, UploadError, ) +from django_scaffold import settings from services.archive import ArchiveService from services.processing.types import UploadArguments from services.test_results import get_flake_set @@ -128,32 +129,43 @@ def process_individual_upload( return False else: flaky_test_set = get_flake_set(db_session, upload.report.commit.repoid) - pg = PGDriver(db_session, flaky_test_set) - bq_enabled = False - if get_config("services", "bigquery", "enabled", default=False): - bq = BQDriver() - bq_enabled = True - - for parsing_info in parsing_infos: - framework = parsing_info["framework"] - testruns = parsing_info["testruns"] - pg.write_testruns( - None, - repoid, - commitid, - branch, - upload, - framework, - testruns, - ) - - if bq_enabled: + + pg = PGDriver(repoid, db_session, flaky_test_set) + if settings.BIGQUERY_WRITE_ENABLED: + bq = BQDriver(repoid) + + for parsing_info in parsing_infos: + framework = parsing_info["framework"] + testruns = parsing_info["testruns"] + pg.write_testruns( + None, + commitid, + branch, + upload.id_, + upload.flag_names, + framework, + testruns, + ) + bq.write_testruns( None, - repoid, commitid, branch, - upload, + upload.id_, + upload.flag_names, + framework, + testruns, + ) + else: + for parsing_info in parsing_infos: + framework = parsing_info["framework"] + testruns = parsing_info["testruns"] + pg.write_testruns( + None, + commitid, + branch, + upload.id_, + upload.flag_names, framework, testruns, ) diff --git a/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json b/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json index 936abe615..83d926a01 100644 --- a/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json +++ b/tasks/tests/unit/snapshots/ta_processor_task__TestUploadTestProcessorTask__ta_processor_task_call__2.json @@ -11,7 +11,6 @@ "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", "framework": "Pytest", "upload_id": "1", - "flags_hash": "AAAAAAAAAAA=", "test_id": "S/K2VdzrrehI4hnoZNsPVg==" }, { @@ -27,7 +26,6 @@ "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", "framework": "Pytest", "upload_id": "1", - "flags_hash": "AAAAAAAAAAA=", "test_id": "CVU2jNUNOkOrl6/lJdK0nw==" }, { @@ -42,7 +40,6 @@ "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", "framework": "Pytest", "upload_id": "1", - "flags_hash": "AAAAAAAAAAA=", "test_id": "UDBibp0NWEToP72TpCn1xg==" }, { @@ -57,7 +54,6 @@ "commit_sha": "cd76b0821854a780b60012aed85af0a8263004ad", "framework": "Pytest", "upload_id": "1", - "flags_hash": "AAAAAAAAAAA=", "test_id": "VE2yD2IYxdSbTvGB6XCJPA==" } ] diff --git a/tasks/tests/unit/test_ta_processor_task.py b/tasks/tests/unit/test_ta_processor_task.py index 9299fa820..4c92d9e50 100644 --- a/tasks/tests/unit/test_ta_processor_task.py +++ b/tasks/tests/unit/test_ta_processor_task.py @@ -33,6 +33,7 @@ def mock_bigquery_service(): class TestUploadTestProcessorTask(object): @pytest.mark.integration @travel("2025-01-01T00:00:00Z", tick=False) + @pytest.mark.django_db(transaction=True, databases=["test_analytics"]) def test_ta_processor_task_call( self, mocker, @@ -45,7 +46,15 @@ def test_ta_processor_task_call( snapshot, ): mock_configuration.set_params( - mock_configuration.params | {"services": {"bigquery": {"enabled": True}}} + mock_configuration.params + | { + "services": {"bigquery": {"write_enabled": True}}, + } + ) + + mocker.patch( + "django_scaffold.settings.BIGQUERY_WRITE_ENABLED", + True, ) tests = dbsession.query(Test).all() diff --git a/worker.sh b/worker.sh index 74ff2f94b..2698c4fae 100755 --- a/worker.sh +++ b/worker.sh @@ -15,6 +15,10 @@ if [ "$RUN_ENV" = "ENTERPRISE" ] || [ "$RUN_ENV" = "DEV" ]; then python manage.py migrate --database "timeseries" fi +if [ "$RUN_ENV" = "DEV" ]; then + python manage.py migrate --database "test_analytics" +fi + if [ -z "$1" ]; then python main.py worker ${queues} From c40a3ae5a939f9798094afd82659d602b1d7f4af Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 17 Jan 2025 14:39:07 -0500 Subject: [PATCH 6/9] create new cache analytics task --- tasks/ta_cache_analytics.py | 81 +++++++++ tasks/tests/unit/test_ta_cache_analytics.py | 184 ++++++++++++++++++++ 2 files changed, 265 insertions(+) create mode 100644 tasks/ta_cache_analytics.py create mode 100644 tasks/tests/unit/test_ta_cache_analytics.py diff --git a/tasks/ta_cache_analytics.py b/tasks/ta_cache_analytics.py new file mode 100644 index 000000000..3b7288ccf --- /dev/null +++ b/tasks/ta_cache_analytics.py @@ -0,0 +1,81 @@ +import datetime as dt + +from redis.exceptions import LockError +from shared.config import get_config +from shared.django_apps.test_analytics.models import LastRollupDate +from shared.utils.enums import TaskConfigGroup +from sqlalchemy.orm import Session + +from app import celery_app +from django_scaffold import settings +from services.redis import get_redis_connection +from ta_storage.bq import BQDriver +from ta_storage.pg import PGDriver +from tasks.base import BaseCodecovTask + +ta_cache_analytics_task_name = ( + f"app.tasks.{TaskConfigGroup.cache_rollup.value}.TACacheAnalyticsTask" +) + + +class TACacheAnalyticsTask(BaseCodecovTask, name=ta_cache_analytics_task_name): + def run_impl( + self, + db_session: Session, + repoid: int, + branch: str, + update_date: bool = True, + **kwargs, + ): + redis_conn = get_redis_connection() + try: + with redis_conn.lock( + f"rollups:{repoid}:{branch}", timeout=300, blocking_timeout=2 + ): + self.run_impl_within_lock(db_session, repoid, branch) + + if update_date: + LastRollupDate.objects.update_or_create( + repoid=repoid, + branch=branch, + defaults={"last_rollup_date": dt.date.today()}, + ) + + with redis_conn.lock(f"rollups:{repoid}", timeout=300, blocking_timeout=2): + self.run_impl_within_lock(db_session, repoid, None) + + if update_date: + LastRollupDate.objects.update_or_create( + repoid=repoid, + branch=None, + defaults={"last_rollup_date": dt.date.today()}, + ) + + except LockError: + return {"in_progress": True} + + return {"success": True} + + def run_impl_within_lock( + self, db_session: Session, repoid: int, branch: str | None + ): + write_buckets = get_config( + "services", "test_analytics", "write_buckets", default=None + ) + + buckets = write_buckets or [ + get_config("services", "minio", "bucket", default="codecov") + ] + + if settings.BIGQUERY_WRITE_ENABLED: + bq = BQDriver(repo_id=repoid) + bq.cache_analytics(buckets, branch) + + pg = PGDriver( + repoid, + ) + pg.cache_analytics(buckets, branch) + + +RegisteredTACacheAnalyticsTask = celery_app.register_task(TACacheAnalyticsTask()) +ta_cache_analytics_task = celery_app.tasks[RegisteredTACacheAnalyticsTask.name] diff --git a/tasks/tests/unit/test_ta_cache_analytics.py b/tasks/tests/unit/test_ta_cache_analytics.py new file mode 100644 index 000000000..7c78dfec0 --- /dev/null +++ b/tasks/tests/unit/test_ta_cache_analytics.py @@ -0,0 +1,184 @@ +import pytest +from redis.exceptions import LockError +from shared.django_apps.core.tests.factories import RepositoryFactory +from shared.django_apps.test_analytics.models import LastRollupDate + +from tasks.ta_cache_analytics import TACacheAnalyticsTask + + +@pytest.fixture +def mock_repo(): + return RepositoryFactory() + + +@pytest.fixture +def mock_redis(mocker): + m = mocker.patch("services.redis._get_redis_instance_from_url") + redis_server = mocker.MagicMock() + m.return_value = redis_server + yield redis_server + + +@pytest.fixture +def mock_config(mock_configuration): + mock_configuration.set_params( + mock_configuration._params + | { + "services": { + "test_analytics": {"write_buckets": ["test-bucket"]}, + "bigquery": {"write_enabled": True}, + } + } + ) + return mock_configuration + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_cache_analytics_disabled_by_settings( + mocker, mock_repo, mock_config, mock_redis +): + # Mock settings to disable BigQuery write + mocker.patch("tasks.ta_cache_analytics.settings.BIGQUERY_WRITE_ENABLED", False) + + mock_bq_driver = mocker.Mock() + mock_pg_driver = mocker.Mock() + mocker.patch("tasks.ta_cache_analytics.BQDriver", return_value=mock_bq_driver) + mocker.patch("tasks.ta_cache_analytics.PGDriver", return_value=mock_pg_driver) + + mock_redis.lock.return_value.__enter__ = lambda x: None + mock_redis.lock.return_value.__exit__ = lambda x, y, z, a: None + + result = TACacheAnalyticsTask().run_impl( + db_session=None, # type: ignore[arg-type] + repoid=mock_repo.repoid, + branch="main", + ) + + assert result == {"success": True} + mock_bq_driver.cache_analytics.assert_not_called() + mock_pg_driver.cache_analytics.assert_has_calls( + [mocker.call(["test-bucket"], "main"), mocker.call(["test-bucket"], None)] + ) + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_cache_analytics_with_bigquery_enabled( + mocker, mock_repo, mock_config, mock_redis +): + mocker.patch("tasks.ta_cache_analytics.settings.BIGQUERY_WRITE_ENABLED", True) + + mock_bq_driver = mocker.Mock() + mock_pg_driver = mocker.Mock() + mocker.patch("tasks.ta_cache_analytics.BQDriver", return_value=mock_bq_driver) + mocker.patch("tasks.ta_cache_analytics.PGDriver", return_value=mock_pg_driver) + + mock_redis.lock.return_value.__enter__ = lambda x: None + mock_redis.lock.return_value.__exit__ = lambda x, y, z, a: None + + result = TACacheAnalyticsTask().run_impl( + db_session=None, # type: ignore[arg-type] + repoid=mock_repo.repoid, + branch="main", + ) + + assert result == {"success": True} + mock_bq_driver.cache_analytics.assert_has_calls( + [mocker.call(["test-bucket"], "main"), mocker.call(["test-bucket"], None)] + ) + mock_pg_driver.cache_analytics.assert_has_calls( + [mocker.call(["test-bucket"], "main"), mocker.call(["test-bucket"], None)] + ) + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_cache_analytics_lock_contention(mocker, mock_repo, mock_config, mock_redis): + mock_redis.lock.side_effect = LockError("Lock already acquired") + + mock_bq_driver = mocker.Mock() + mock_pg_driver = mocker.Mock() + mocker.patch("tasks.ta_cache_analytics.BQDriver", return_value=mock_bq_driver) + mocker.patch("tasks.ta_cache_analytics.PGDriver", return_value=mock_pg_driver) + + result = TACacheAnalyticsTask().run_impl( + db_session=None, # type: ignore[arg-type] + repoid=mock_repo.repoid, + branch="main", + ) + + assert result == {"in_progress": True} + mock_bq_driver.cache_analytics.assert_not_called() + mock_pg_driver.cache_analytics.assert_not_called() + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_cache_analytics_updates_last_rollup_date( + mocker, mock_repo, mock_config, mock_redis +): + mocker.patch("tasks.ta_cache_analytics.settings.BIGQUERY_WRITE_ENABLED", True) + + mock_bq_driver = mocker.Mock() + mock_pg_driver = mocker.Mock() + mocker.patch("tasks.ta_cache_analytics.BQDriver", return_value=mock_bq_driver) + mocker.patch("tasks.ta_cache_analytics.PGDriver", return_value=mock_pg_driver) + + mock_redis.lock.return_value.__enter__ = lambda x: None + mock_redis.lock.return_value.__exit__ = lambda x, y, z, a: None + + result = TACacheAnalyticsTask().run_impl( + db_session=None, # type: ignore[arg-type] + repoid=mock_repo.repoid, + branch="main", + ) + + assert result == {"success": True} + + # Verify LastRollupDate was updated for both branch and repo-level + branch_rollup = LastRollupDate.objects.get(repoid=mock_repo.repoid, branch="main") + assert branch_rollup is not None + + repo_rollup = LastRollupDate.objects.get(repoid=mock_repo.repoid, branch=None) + assert repo_rollup is not None + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_cache_analytics_with_custom_buckets( + mocker, mock_repo, mock_redis, mock_config +): + mock_config.set_params( + mock_config._params + | { + "services": { + "test_analytics": {"write_buckets": ["bucket1", "bucket2"]}, + "bigquery": {"write_enabled": True}, + } + } + ) + mocker.patch("tasks.ta_cache_analytics.settings.BIGQUERY_WRITE_ENABLED", True) + + mock_bq_driver = mocker.Mock() + mock_pg_driver = mocker.Mock() + mocker.patch("tasks.ta_cache_analytics.BQDriver", return_value=mock_bq_driver) + mocker.patch("tasks.ta_cache_analytics.PGDriver", return_value=mock_pg_driver) + + mock_redis.lock.return_value.__enter__ = lambda x: None + mock_redis.lock.return_value.__exit__ = lambda x, y, z, a: None + + result = TACacheAnalyticsTask().run_impl( + db_session=None, # type: ignore[arg-type] + repoid=mock_repo.repoid, + branch="main", + ) + + assert result == {"success": True} + mock_bq_driver.cache_analytics.assert_has_calls( + [ + mocker.call(["bucket1", "bucket2"], "main"), + mocker.call(["bucket1", "bucket2"], None), + ] + ) + mock_pg_driver.cache_analytics.assert_has_calls( + [ + mocker.call(["bucket1", "bucket2"], "main"), + mocker.call(["bucket1", "bucket2"], None), + ] + ) From 4f1273f683d6ec81705d7bc593b4151e727d7453 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 17 Jan 2025 14:39:07 -0500 Subject: [PATCH 7/9] create new process flakes task --- tasks/ta_process_flakes.py | 81 +++++++++ tasks/tests/unit/test_ta_process_flakes.py | 196 +++++++++++++++++++++ 2 files changed, 277 insertions(+) create mode 100644 tasks/ta_process_flakes.py create mode 100644 tasks/tests/unit/test_ta_process_flakes.py diff --git a/tasks/ta_process_flakes.py b/tasks/ta_process_flakes.py new file mode 100644 index 000000000..156818554 --- /dev/null +++ b/tasks/ta_process_flakes.py @@ -0,0 +1,81 @@ +import logging +from typing import Any + +from redis.exceptions import LockError +from shared.django_apps.reports.models import CommitReport, ReportSession +from shared.utils.enums import TaskConfigGroup + +from app import celery_app +from django_scaffold import settings +from services.redis import get_redis_connection +from ta_storage.bq import BQDriver +from ta_storage.pg import PGDriver +from tasks.base import BaseCodecovTask + +log = logging.getLogger(__name__) + +TA_FLAKE_LOCK_KEY = "ta_flake_lock:{repo_id}" +TA_FLAKE_UPLOADS_KEY = "ta_flake_uploads:{repo_id}" + +FLAKE_EXPIRY_COUNT = 30 + +ta_process_flakes_task_name = ( + f"app.tasks.{TaskConfigGroup.flakes.value}.TAProcessFlakesTask" +) + + +class TAProcessFlakesTask(BaseCodecovTask, name=ta_process_flakes_task_name): + """ + This task is currently called in the test results finisher task and in the sync pulls task + """ + + def run_impl( + self, + _db_session: Any, + *, + repo_id: int, + commit_id: str, + **kwargs: Any, + ): + log.info( + "Received process flakes task", + extra=dict(repoid=repo_id, commit=commit_id), + ) + + redis_client = get_redis_connection() + lock_name = f"ta_flake_lock:{repo_id}" + try: + with redis_client.lock( + lock_name, timeout=max(300, self.hard_time_limit_task), blocking=False + ): + while redis_client.get(f"ta_flake_uploads:{repo_id}") is not None: + redis_client.delete(f"ta_flake_uploads:{repo_id}") + process_flakes_for_repo(repo_id, commit_id) + except LockError: + log.warning("Unable to acquire process flakeslock for key %s.", lock_name) + return {"successful": False} + + return {"successful": True} + + +def process_flakes_for_repo(repo_id: int, commit_id): + # get all uploads pending process flakes in the entire repo? why stop at a given commit :D + uploads_to_process = ReportSession.objects.filter( + report__report_type=CommitReport.ReportType.TEST_RESULTS.value, + report__commit__repository__repoid=repo_id, + report__commit__commitid=commit_id, + state__in=["v2_finished"], + ).all() + if not uploads_to_process: + return + + if settings.BIGQUERY_WRITE_ENABLED: + bq = BQDriver(repo_id) + bq.write_flakes([upload for upload in uploads_to_process]) + + pg = PGDriver(repo_id) + pg.write_flakes([upload for upload in uploads_to_process]) + + +TAProcessFlakesTaskRegistered = celery_app.register_task(TAProcessFlakesTask()) +ta_process_flakes_task = celery_app.tasks[TAProcessFlakesTaskRegistered.name] diff --git a/tasks/tests/unit/test_ta_process_flakes.py b/tasks/tests/unit/test_ta_process_flakes.py new file mode 100644 index 000000000..73b69ed04 --- /dev/null +++ b/tasks/tests/unit/test_ta_process_flakes.py @@ -0,0 +1,196 @@ +import pytest +from redis.exceptions import LockError +from shared.config import ConfigHelper +from shared.django_apps.core.tests.factories import CommitFactory, RepositoryFactory +from shared.django_apps.reports.models import CommitReport +from shared.django_apps.reports.tests.factories import UploadFactory + +from tasks.ta_process_flakes import TAProcessFlakesTask + + +@pytest.fixture +def mock_repo(): + return RepositoryFactory() + + +@pytest.fixture +def mock_redis(mocker): + m = mocker.patch("services.redis._get_redis_instance_from_url") + redis_server = mocker.MagicMock() + m.return_value = redis_server + yield redis_server + + +@pytest.fixture +def mock_config(mock_configuration): + mock_config = ConfigHelper() + mock_config.set_params( + mock_configuration._params + | { + "setup": {"test_analytics_database": {"enabled": True}}, + "services": { + "bigquery": { + "write_enabled": True, + "read_enabled": True, + } + }, + } + ) + return mock_config + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_process_flakes_disabled_by_config( + mocker, mock_repo, mock_configuration, mock_redis +): + mock_config = ConfigHelper() + mock_config.set_params( + mock_configuration._params + | {"setup": {"test_analytics_database": {"enabled": False}}} + ) + mocker.patch("django_scaffold.settings.get_config", return_value=mock_config) + + mock_driver = mocker.Mock() + mock_driver_cls = mocker.patch( + "tasks.ta_process_flakes.BQDriver", return_value=mock_driver + ) + + mock_redis.get.side_effect = ["1", None] + + commit = CommitFactory(repository=mock_repo) + result = TAProcessFlakesTask().run_impl( + _db_session=None, + repo_id=mock_repo.repoid, + commit_id=commit.commitid, + ) + + assert result == {"successful": True} + mock_driver_cls.assert_not_called() + mock_driver.write_flakes.assert_not_called() + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_process_flakes_disabled_by_settings( + mocker, mock_repo, mock_config, mock_redis +): + # Mock settings to disable BigQuery write + mocker.patch("tasks.ta_process_flakes.settings.BIGQUERY_WRITE_ENABLED", False) + + mock_driver = mocker.Mock() + mock_driver_cls = mocker.patch( + "tasks.ta_process_flakes.BQDriver", return_value=mock_driver + ) + + mock_redis.get.side_effect = ["1", None] + + commit = CommitFactory(repository=mock_repo) + result = TAProcessFlakesTask().run_impl( + _db_session=None, + repo_id=mock_repo.repoid, + commit_id=commit.commitid, + ) + + assert result == {"successful": True} + mock_driver_cls.assert_not_called() + mock_driver.write_flakes.assert_not_called() + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_process_flakes_no_uploads(mocker, mock_repo, mock_config, mock_redis): + # Mock settings to enable BigQuery write + mocker.patch("tasks.ta_process_flakes.settings.BIGQUERY_WRITE_ENABLED", True) + + mock_driver = mocker.Mock() + mock_driver_cls = mocker.patch( + "tasks.ta_process_flakes.BQDriver", return_value=mock_driver + ) + + mock_redis.get.side_effect = ["1", None] + + commit = CommitFactory(repository=mock_repo) + result = TAProcessFlakesTask().run_impl( + _db_session=None, + repo_id=mock_repo.repoid, + commit_id=commit.commitid, + ) + + assert result == {"successful": True} + mock_driver_cls.assert_not_called() + mock_driver.write_flakes.assert_not_called() + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_process_flakes_with_uploads(mocker, mock_repo, mock_config, mock_redis): + # Mock settings to enable BigQuery write + mocker.patch("tasks.ta_process_flakes.settings.BIGQUERY_WRITE_ENABLED", True) + + mock_driver = mocker.Mock() + mock_driver_cls = mocker.patch( + "tasks.ta_process_flakes.BQDriver", return_value=mock_driver + ) + + # Configure redis.get to return value first then None + mock_redis.get.side_effect = ["1", None] + + commit = CommitFactory(repository=mock_repo) + + # Create multiple uploads with different states + upload1 = UploadFactory( + report__report_type=CommitReport.ReportType.TEST_RESULTS.value, + report__commit=commit, + state="v2_finished", + ) + upload2 = UploadFactory( + report__report_type=CommitReport.ReportType.TEST_RESULTS.value, + report__commit=commit, + state="v2_finished", + ) + # Upload that should be ignored due to state + UploadFactory( + report__report_type=CommitReport.ReportType.TEST_RESULTS.value, + report__commit=commit, + state="processing", + ) + # Upload that should be ignored due to report type + UploadFactory( + report__report_type=CommitReport.ReportType.COVERAGE.value, + report__commit=commit, + state="v2_finished", + ) + + result = TAProcessFlakesTask().run_impl( + _db_session=None, + repo_id=mock_repo.repoid, + commit_id=commit.commitid, + ) + + assert result == {"successful": True} + mock_driver_cls.assert_called_once_with(mock_repo.repoid) + mock_driver.write_flakes.assert_called_once() + # Verify the uploads passed to write_flakes + uploads_processed = mock_driver.write_flakes.call_args[0][0] + assert len(uploads_processed) == 2 + assert set(u.id for u in uploads_processed) == {upload1.id, upload2.id} + # Verify redis.get was called twice + assert mock_redis.get.call_count == 2 + + +@pytest.mark.django_db(transaction=True, databases={"default", "test_analytics"}) +def test_ta_process_flakes_lock_contention(mocker, mock_repo, mock_config, mock_redis): + # Mock settings to enable BigQuery write + mocker.patch("tasks.ta_process_flakes.settings.BIGQUERY_WRITE_ENABLED", True) + + mock_driver = mocker.Mock() + mocker.patch("tasks.ta_process_flakes.BQDriver", return_value=mock_driver) + + mock_redis.lock = mocker.Mock(side_effect=LockError("Lock already acquired")) + + commit = CommitFactory(repository=mock_repo) + result = TAProcessFlakesTask().run_impl( + _db_session=None, + repo_id=mock_repo.repoid, + commit_id=commit.commitid, + ) + + assert result == {"successful": False} + mock_driver.write_flakes.assert_not_called() From 6130eca672bd4e6127708d4f32d424fbab4748b9 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 17 Jan 2025 14:39:07 -0500 Subject: [PATCH 8/9] consume new TADriver and flake, cache tasks in finisher --- tasks/ta_finisher.py | 277 +++++++++++++--------- tasks/tests/unit/test_ta_finisher_task.py | 41 +++- 2 files changed, 200 insertions(+), 118 deletions(-) diff --git a/tasks/ta_finisher.py b/tasks/ta_finisher.py index c5d5d2e4e..ad53bc1e5 100644 --- a/tasks/ta_finisher.py +++ b/tasks/ta_finisher.py @@ -13,12 +13,12 @@ from database.models import ( Commit, CommitReport, - Flake, Repository, TestResultReportTotals, Upload, UploadError, ) +from django_scaffold.settings import BIGQUERY_READ_ENABLED from helpers.checkpoint_logger.flows import TestResultsFlow from helpers.notifier import NotifierResult from helpers.string import EscapeEnum, Replacement, StringEscaper, shorten_file_paths @@ -32,20 +32,20 @@ get_repo_provider_service, ) from services.seats import ShouldActivateSeat, determine_seat_activation -from services.test_results import ( +from services.ta_utils import ( FlakeInfo, - TACommentInDepthInfo, - TestResultsNotificationFailure, + TestFailure, TestResultsNotificationPayload, TestResultsNotifier, - get_test_summary_for_commit, - latest_failures_for_commit, should_do_flaky_detection, ) +from ta_storage.base import PRCommentAggResult, PRCommentFailResult +from ta_storage.bq import BQDriver +from ta_storage.pg import PGDriver from tasks.base import BaseCodecovTask -from tasks.cache_test_rollups import cache_test_rollups_task from tasks.notify import notify_task -from tasks.process_flakes import process_flakes_task +from tasks.ta_cache_analytics import ta_cache_analytics_task +from tasks.ta_process_flakes import TA_FLAKE_UPLOADS_KEY, ta_process_flakes_task log = logging.getLogger(__name__) @@ -86,18 +86,18 @@ def queue_optional_tasks( branch: str | None, ): redis_client = get_redis_connection() - if should_do_flaky_detection(repo, commit_yaml): if commit.merged is True or branch == repo.branch: - redis_client.set(f"flake_uploads:{repo.repoid}", 0) - process_flakes_task_sig = process_flakes_task.s( + # run new process flakes task + redis_client.set(TA_FLAKE_UPLOADS_KEY.format(repo_id=repo.repoid), 0) + ta_process_flakes_task_sig = ta_process_flakes_task.s( repo_id=repo.repoid, commit_id=commit.commitid, ) - process_flakes_task_sig.apply_async() + ta_process_flakes_task_sig.apply_async() if branch is not None: - cache_task_sig = cache_test_rollups_task.s( + cache_task_sig = ta_cache_analytics_task.s( repoid=repo.repoid, branch=branch, ) @@ -109,9 +109,8 @@ def get_totals( ) -> TestResultReportTotals: totals = commit_report.test_result_totals if totals is None: - totals = TestResultReportTotals( - report_id=commit_report.id, - ) + totals = TestResultReportTotals() + totals.report = commit_report totals.passed = 0 totals.skipped = 0 totals.failed = 0 @@ -121,64 +120,42 @@ def get_totals( return totals -def populate_failures( - failures: list[TestResultsNotificationFailure], - db_session: Session, - repoid: int, - commitid: str, - shorten_paths: bool, - uploads: dict[int, Upload], - escaper: StringEscaper, -) -> None: - failed_test_instances = latest_failures_for_commit(db_session, repoid, commitid) - - for test_instance in failed_test_instances: - failure_message = test_instance.failure_message - if failure_message is not None: - if shorten_paths: - failure_message = shorten_file_paths(failure_message) - failure_message = escaper.replace(failure_message) - - upload = uploads[test_instance.upload_id] - - failures.append( - TestResultsNotificationFailure( - display_name=test_instance.test.computed_name - if test_instance.test.computed_name is not None - else test_instance.test.name, - failure_message=failure_message, - test_id=test_instance.test_id, - envs=upload.flag_names, - duration_seconds=test_instance.duration_seconds, - build_url=upload.build_url, - ) +def get_bigquery_test_data( + repo: Repository, commit_sha: str, commit_yaml: UserYaml +) -> tuple[ + PRCommentAggResult, + list[PRCommentFailResult[tuple[bytes, bytes | None]]], + dict[tuple[bytes, bytes | None], FlakeInfo] | None, +]: + driver = BQDriver(repo.repoid) + agg_result = driver.pr_comment_agg(commit_sha) + failures = driver.pr_comment_fail(commit_sha) + if should_do_flaky_detection(repo, commit_yaml): + flaky_tests = driver.get_repo_flakes( + tuple(failure["id"] for failure in failures) ) + else: + flaky_tests = None + return agg_result, failures, flaky_tests -def get_flaky_tests( - db_session: Session, - repoid: int, - failures: list[TestResultsNotificationFailure], -) -> dict[str, FlakeInfo]: - failure_test_ids = [failure.test_id for failure in failures] - - matching_flakes = list( - db_session.query(Flake) - .filter( - Flake.repoid == repoid, - Flake.testid.in_(failure_test_ids), - Flake.end_date.is_(None), - Flake.count != (Flake.recent_passes_count + Flake.fail_count), + +def get_postgres_test_data( + db_session: Session, repo: Repository, commit_sha: str, commit_yaml: UserYaml +) -> tuple[ + PRCommentAggResult, list[PRCommentFailResult[str]], dict[str, FlakeInfo] | None +]: + driver = PGDriver(repo.repoid, db_session) + agg_result = driver.pr_comment_agg(commit_sha) + failures = driver.pr_comment_fail(commit_sha) + if should_do_flaky_detection(repo, commit_yaml): + flaky_tests = driver.get_repo_flakes( + tuple(failure["id"] for failure in failures) ) - .limit(100) - .all() - ) + else: + flaky_tests = None - flaky_test_ids = { - flake.testid: FlakeInfo(flake.fail_count, flake.count) - for flake in matching_flakes - } - return flaky_test_ids + return agg_result, failures, flaky_tests class TAFinisherTask(BaseCodecovTask, name=ta_finisher_task_name): @@ -243,59 +220,143 @@ def process_impl_within_lock( log.info("Running test results finishers", extra=self.extra_dict) TestResultsFlow.log(TestResultsFlow.TEST_RESULTS_FINISHER_BEGIN) - commit: Commit = ( + commit = ( db_session.query(Commit).filter_by(repoid=repoid, commitid=commitid).first() ) - assert commit, "commit not found" + if commit is None: + raise ValueError("commit not found") commit_report = commit.commit_report(ReportType.TEST_RESULTS) - totals = get_totals(commit_report, db_session) - uploads = get_uploads( - db_session, commit - ) # processed uploads that have yet to be persisted + uploads = get_uploads(db_session, commit) repo = commit.repository branch = commit.branch - uploads = get_uploads(db_session, commit) + if BIGQUERY_READ_ENABLED: + agg_result, failures, flaky_tests = get_bigquery_test_data( + repo, commitid, commit_yaml + ) + payload = TestResultsNotificationPayload( + agg_result["failed_ct"] + agg_result["flaky_failed_ct"], + agg_result["passed_ct"], + agg_result["skipped_ct"], + ) - # if we succeed once, error should be None for this commit forever - if totals.error is not None: - totals.error = None - db_session.flush() + if failures: + escaper = StringEscaper(ESCAPE_FAILURE_MESSAGE_DEFN) + shorten_paths = commit_yaml.read_yaml_field( + "test_analytics", "shorten_paths", _else=True + ) - test_summary = get_test_summary_for_commit(db_session, repoid, commitid) - totals.failed = test_summary.get("error", 0) + test_summary.get("failure", 0) - totals.skipped = test_summary.get("skip", 0) - totals.passed = test_summary.get("pass", 0) - db_session.flush() + failures_list = [] + flaky_failures = [] + + for failure in failures: + failure_message = failure["failure_message"] + if failure_message is not None and shorten_paths: + failure_message = shorten_file_paths(failure_message) + if failure_message is not None: + failure_message = escaper.replace(failure_message) + + test_id = failure["id"] + base_failure = TestFailure( + display_name=failure["computed_name"], + failure_message=failure_message, + duration_seconds=failure["duration_seconds"], + build_url=uploads[failure["upload_id"]].build_url + if failure["upload_id"] in uploads + else None, + ) + + if flaky_tests and test_id in flaky_tests: + flaky_failures.append( + TestFailure( + display_name=base_failure.display_name, + failure_message=base_failure.failure_message, + duration_seconds=base_failure.duration_seconds, + build_url=base_failure.build_url, + flake_info=flaky_tests[test_id], + ) + ) + else: + failures_list.append(base_failure) + + failures_list = sorted(failures_list, key=lambda x: x.duration_seconds) + flaky_failures = sorted( + flaky_failures, key=lambda x: x.duration_seconds + ) + + payload.regular_failures = failures_list if failures_list else None + payload.flaky_failures = flaky_failures if flaky_failures else None + else: + totals = get_totals(commit_report, db_session) - info = None - if totals.failed: - escaper = StringEscaper(ESCAPE_FAILURE_MESSAGE_DEFN) - shorten_paths = commit_yaml.read_yaml_field( - "test_analytics", "shorten_paths", _else=True + agg_result, failures, flaky_tests = get_postgres_test_data( + db_session, repo, commitid, commit_yaml ) - failures = [] - populate_failures( - failures, - db_session, - repoid, - commitid, - shorten_paths, - uploads, - escaper, + totals.failed = agg_result["failed_ct"] + totals.skipped = agg_result["skipped_ct"] + totals.passed = agg_result["passed_ct"] + db_session.flush() + + payload = TestResultsNotificationPayload( + totals.failed, totals.passed, totals.skipped ) - flaky_tests = dict() - if should_do_flaky_detection(repo, commit_yaml): - flaky_tests = get_flaky_tests(db_session, repoid, failures) + if failures: + escaper = StringEscaper(ESCAPE_FAILURE_MESSAGE_DEFN) + shorten_paths = commit_yaml.read_yaml_field( + "test_analytics", "shorten_paths", _else=True + ) - failures = sorted(failures, key=lambda x: x.duration_seconds)[:3] + failures_list = [] + flaky_failures = [] + + for failure in failures: + failure_message = failure["failure_message"] + if failure_message is not None and shorten_paths: + failure_message = shorten_file_paths(failure_message) + if failure_message is not None: + failure_message = escaper.replace(failure_message) + + test_id = failure["id"] + base_failure = TestFailure( + display_name=failure["computed_name"], + failure_message=failure_message, + duration_seconds=failure["duration_seconds"], + build_url=uploads[failure["upload_id"]].build_url + if failure["upload_id"] in uploads + else None, + ) + + if flaky_tests and test_id in flaky_tests: + flaky_failures.append( + TestFailure( + display_name=base_failure.display_name, + failure_message=base_failure.failure_message, + duration_seconds=base_failure.duration_seconds, + build_url=base_failure.build_url, + flake_info=flaky_tests[test_id], + ) + ) + else: + failures_list.append(base_failure) + + failures_list = sorted(failures_list, key=lambda x: x.duration_seconds) + flaky_failures = sorted( + flaky_failures, key=lambda x: x.duration_seconds + ) - info = TACommentInDepthInfo(failures, flaky_tests) + if failures_list or flaky_failures: + payload = TestResultsNotificationPayload( + totals.failed, + totals.passed, + totals.skipped, + regular_failures=failures_list if failures_list else None, + flaky_failures=flaky_failures if flaky_failures else None, + ) additional_data: AdditionalData = {"upload_type": UploadType.TEST_RESULTS} repo_service = get_repo_provider_service(repo, additional_data=additional_data) @@ -316,16 +377,13 @@ def process_impl_within_lock( .first() ) - if not (info or upload_error): + if not (payload or upload_error): return { "notify_attempted": False, "notify_succeeded": False, "queue_notify": True, } - payload = TestResultsNotificationPayload( - totals.failed, totals.passed, totals.skipped, info - ) notifier = TestResultsNotifier( commit, commit_yaml, @@ -335,6 +393,7 @@ def process_impl_within_lock( error=upload_error, ) notifier_result = notifier.notify() + for upload in uploads.values(): upload.state = "v2_finished" db_session.commit() @@ -351,7 +410,7 @@ def process_impl_within_lock( return { "notify_attempted": True, "notify_succeeded": success, - "queue_notify": not info, + "queue_notify": not payload, } def seat_activation( diff --git a/tasks/tests/unit/test_ta_finisher_task.py b/tasks/tests/unit/test_ta_finisher_task.py index e4464bdee..1f4d1c338 100644 --- a/tasks/tests/unit/test_ta_finisher_task.py +++ b/tasks/tests/unit/test_ta_finisher_task.py @@ -136,6 +136,7 @@ def test_test_analytics( owner__username="joseph-sentry", owner__service="github", name="codecov-demo", + branch="main", ) dbsession.add(repo) dbsession.flush() @@ -162,10 +163,15 @@ def test_test_analytics( mocker.patch.object(TAProcessorTask, "app", celery_app) mocker.patch.object(TAFinisherTask, "app", celery_app) - celery_app.tasks = { - "app.tasks.flakes.ProcessFlakesTask": mocker.MagicMock(), - "app.tasks.cache_rollup.CacheTestRollupsTask": mocker.MagicMock(), - } + mock_flakes_sig = mocker.MagicMock() + mock_flakes_sig.apply_async.return_value = None + mock_flakes_task = mocker.patch("tasks.ta_finisher.ta_process_flakes_task") + mock_flakes_task.s.return_value = mock_flakes_sig + + mock_cache_sig = mocker.MagicMock() + mock_cache_sig.apply_async.return_value = None + mock_cache_task = mocker.patch("tasks.ta_finisher.ta_cache_analytics_task") + mock_cache_task.s.return_value = mock_cache_sig pull = PullFactory.create(repository=commit.repository, head=commit.commitid) dbsession.add(pull) @@ -256,12 +262,29 @@ def test_test_analytics( commit_yaml={"codecov": {"max_report_age": False}}, ) - assert result["notify_attempted"] is True - assert result["notify_succeeded"] is True - assert result["queue_notify"] is False + expected_result = { + "notify_attempted": True, + "notify_succeeded": True, + "queue_notify": False, + } + + assert result == expected_result mock_repo_provider_comments.post_comment.assert_called_once() + # Verify task signatures were called correctly + mock_flakes_task.s.assert_called_once_with( + repo_id=repo.repoid, + commit_id=commit.commitid, + ) + mock_flakes_sig.apply_async.assert_called_once() + + mock_cache_task.s.assert_called_once_with( + repoid=repo.repoid, + branch=commit.branch, + ) + mock_cache_sig.apply_async.assert_called_once() + short_form_service_name = services_short_dict.get( upload.report.commit.repository.owner.service ) @@ -276,7 +299,7 @@ def test_test_analytics( > > ```python -> tests.test_parsers.TestParsers test_divide +> hello_world.tests.test_parsers.TestParsers.test_divide > ``` > >
Stack Traces | 0.001s run time @@ -291,7 +314,7 @@ def test_test_analytics( > > ```python -> tests.test_parsers.TestParsers test_subtract +> hello_world.tests.test_parsers.TestParsers.test_subtract > ``` > >
Stack Traces | 0.004s run time From 7bcdb96accd37b9800e86c9e0915f637495f6e01 Mon Sep 17 00:00:00 2001 From: Joseph Sawaya Date: Fri, 17 Jan 2025 14:39:07 -0500 Subject: [PATCH 9/9] quick fix in ta processor --- tasks/ta_processor.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tasks/ta_processor.py b/tasks/ta_processor.py index 64ed643b1..cce013000 100644 --- a/tasks/ta_processor.py +++ b/tasks/ta_processor.py @@ -17,7 +17,6 @@ from django_scaffold import settings from services.archive import ArchiveService from services.processing.types import UploadArguments -from services.test_results import get_flake_set from services.yaml import read_yaml_field from ta_storage.bq import BQDriver from ta_storage.pg import PGDriver @@ -128,9 +127,8 @@ def process_individual_upload( db_session.commit() return False else: - flaky_test_set = get_flake_set(db_session, upload.report.commit.repoid) - - pg = PGDriver(repoid, db_session, flaky_test_set) + # the flaky test set will be generated in the first call to write_testruns + pg = PGDriver(repoid, db_session) if settings.BIGQUERY_WRITE_ENABLED: bq = BQDriver(repoid)