fix: address feedback

generated_proto/testrun/ta_testrun_pb2.pyi

@@ -0,0 +1,83 @@
+from typing import ClassVar as _ClassVar
+from typing import Iterable as _Iterable
+from typing import Optional as _Optional
+from typing import Union as _Union
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class TestRun(_message.Message):
+    __slots__ = (
+        "timestamp",
+        "name",
+        "classname",
+        "testsuite",
+        "computed_name",
+        "outcome",
+        "failure_message",
+        "duration_seconds",
+        "repoid",
+        "commit_sha",
+        "branch_name",
+        "flags",
+        "filename",
+        "framework",
+    )
+    class Outcome(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = ()
+        PASSED: _ClassVar[TestRun.Outcome]
+        FAILED: _ClassVar[TestRun.Outcome]
+        SKIPPED: _ClassVar[TestRun.Outcome]
+
+    PASSED: TestRun.Outcome
+    FAILED: TestRun.Outcome
+    SKIPPED: TestRun.Outcome
+    TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    CLASSNAME_FIELD_NUMBER: _ClassVar[int]
+    TESTSUITE_FIELD_NUMBER: _ClassVar[int]
+    COMPUTED_NAME_FIELD_NUMBER: _ClassVar[int]
+    OUTCOME_FIELD_NUMBER: _ClassVar[int]
+    FAILURE_MESSAGE_FIELD_NUMBER: _ClassVar[int]
+    DURATION_SECONDS_FIELD_NUMBER: _ClassVar[int]
+    REPOID_FIELD_NUMBER: _ClassVar[int]
+    COMMIT_SHA_FIELD_NUMBER: _ClassVar[int]
+    BRANCH_NAME_FIELD_NUMBER: _ClassVar[int]
+    FLAGS_FIELD_NUMBER: _ClassVar[int]
+    FILENAME_FIELD_NUMBER: _ClassVar[int]
+    FRAMEWORK_FIELD_NUMBER: _ClassVar[int]
+    timestamp: int
+    name: str
+    classname: str
+    testsuite: str
+    computed_name: str
+    outcome: TestRun.Outcome
+    failure_message: str
+    duration_seconds: float
+    repoid: int
+    commit_sha: str
+    branch_name: str
+    flags: _containers.RepeatedScalarFieldContainer[str]
+    filename: str
+    framework: str
+    def __init__(
+        self,
+        timestamp: _Optional[int] = ...,
+        name: _Optional[str] = ...,
+        classname: _Optional[str] = ...,
+        testsuite: _Optional[str] = ...,
+        computed_name: _Optional[str] = ...,
+        outcome: _Optional[_Union[TestRun.Outcome, str]] = ...,
+        failure_message: _Optional[str] = ...,
+        duration_seconds: _Optional[float] = ...,
+        repoid: _Optional[int] = ...,
+        commit_sha: _Optional[str] = ...,
+        branch_name: _Optional[str] = ...,
+        flags: _Optional[_Iterable[str]] = ...,
+        filename: _Optional[str] = ...,
+        framework: _Optional[str] = ...,
+    ) -> None: ...

protobuf/ta_testrun.proto

@@ -1,7 +1,7 @@
 syntax = "proto2";
 
 message TestRun {
-    optional string timestamp = 1;
+    optional int64 timestamp = 1;
     optional string name = 2;
     optional string classname = 3;
     optional string testsuite = 4;
@@ -16,12 +16,12 @@ message TestRun {
     optional Outcome outcome = 6;
 
     optional string failure_message = 7;
-    optional float duration = 8;
+    optional float duration_seconds = 8;
 
     optional int64 repoid = 10;
     optional string commit_sha = 11;
 
-    optional string branch = 12;
+    optional string branch_name = 12;
 
     repeated string flags = 13;
 

requirements.in

@@ -28,7 +28,7 @@ pre-commit
 polars==1.12.0
 proto-plus>=1.25.0
 psycopg2>=2.9.10
-protobuf>=5.29.2`
+protobuf>=5.29.2
 pydantic>=2.9.0
 PyJWT>=2.4.0
 pytest

services/tests/test_bigquery.py

@@ -4,7 +4,7 @@
 import polars as pl
 import pytest
 
-import generated_proto.sample.sample_pb2 as sample_pb2
+import generated_proto.testrun.ta_testrun_pb2 as ta_testrun_pb2
 from services.bigquery import BigQueryService
 
 fake_private_key = """-----BEGIN RSA PRIVATE KEY-----
@@ -55,12 +55,12 @@ def test_bigquery_service():
     results = bigquery_service.query(sql)
 
     assert len(results) == 2
-    assert set([row["timestamp"] for row in results]) == {
+    assert {row["timestamp"] for row in results} == {
         datetime.fromisoformat("2025-01-01T00:00:00Z"),
         datetime.fromisoformat("2024-12-30T00:00:00Z"),
     }
-    assert set([row["name"] for row in results]) == {"name", "name2"}
-    assert set([row["id"] for row in results]) == {1, 2}
+    assert {row["name"] for row in results} == {"name", "name2"}
+    assert {row["id"] for row in results} == {1, 2}
 
 
 @pytest.mark.skip(reason="This test requires being run using actual working creds")
@@ -78,53 +78,71 @@ def test_bigquery_service_polars():
     )
 
     assert len(results) == 2
-    assert set(results["timestamp"].to_list()) == {
+    assert {x for x in results["timestamp"].to_list()} == {
         datetime.fromisoformat("2025-01-01T00:00:00Z"),
         datetime.fromisoformat("2024-12-30T00:00:00Z"),
     }
-    assert set(results["name"].to_list()) == {"name", "name2"}
-    assert set(results["id"].to_list()) == {1, 2}
+    assert {x for x in results["name"].to_list()} == {"name", "name2"}
+    assert {x for x in results["id"].to_list()} == {1, 2}
 
 
 # this test should only be run manually when making changes to the way we write to bigquery
 # the reason it's not automated is because vcrpy does not seem to work with the gRPC requests
 @pytest.mark.skip(reason="This test requires being run using actual working creds")
 def test_bigquery_service_write():
-    table_name = "codecov-dev.test_dataset.sample_table"
+    table_name = "codecov-dev.test_dataset.testruns"
     bigquery_service = BigQueryService(gcp_config)
 
     bigquery_service.query(f"TRUNCATE TABLE `{table_name}`")
 
     data = [
-        sample_pb2.Sample(
-            timestamp="2025-01-01T00:00:00Z",
-            id="1",
+        ta_testrun_pb2.TestRun(
+            timestamp=int(
+                datetime.fromisoformat("2025-01-01T00:00:00.000000Z").timestamp()
+                * 1000000
+            ),
             name="name",
+            classname="classname",
+            testsuite="testsuite",
+            computed_name="computed_name",
+            outcome=ta_testrun_pb2.TestRun.Outcome.PASSED,
+            failure_message="failure_message",
+            duration_seconds=1.0,
+            filename="filename",
         ),
-        sample_pb2.Sample(
-            timestamp="2024-12-30T00:00:00Z",
-            id="2",
+        ta_testrun_pb2.TestRun(
+            timestamp=int(
+                datetime.fromisoformat("2024-12-30T00:00:00.000000Z").timestamp()
+                * 1000000
+            ),
             name="name2",
+            classname="classname2",
+            testsuite="testsuite2",
+            computed_name="computed_name2",
+            outcome=ta_testrun_pb2.TestRun.Outcome.FAILED,
+            failure_message="failure_message2",
+            duration_seconds=2.0,
+            filename="filename2",
         ),
     ]
 
     serialized_data = [row.SerializeToString() for row in data]
 
     bigquery_service.write(
         "test_dataset",
-        "test_table",
-        sample_pb2,
+        "testruns",
+        ta_testrun_pb2,
         serialized_data,
     )
 
     results = bigquery_service.query(f"SELECT * FROM `{table_name}`")
 
     assert len(results) == 2
 
-    assert set([row["timestamp"] for row in results]) == set(
+    assert {row["timestamp"] for row in results} == set(
         [
             datetime.fromisoformat("2025-01-01T00:00:00Z"),
             datetime.fromisoformat("2024-12-30T00:00:00Z"),
         ]
     )
-    assert set([row["name"] for row in results]) == set(["name", "name2"])
+    assert {row["name"] for row in results} == set(["name", "name2"])

ta_storage/base.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
-from typing import Any
+
+from test_results_parser import Testrun
 
 from database.models.reports import Upload
 
@@ -8,11 +9,12 @@ class TADriver(ABC):
     @abstractmethod
     def write_testruns(
         self,
+        timestamp: int,
         repo_id: int,
-        commit_id: str,
-        branch: str,
+        commit_sha: str,
+        branch_name: str,
         upload: Upload,
         framework: str | None,
-        testruns: list[dict[str, Any]],
+        testruns: list[Testrun],
     ):
         pass