Tests

Author: spanglerco

tests/integration/dataframe/test_dataframe.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
+import warnings
 from datetime import datetime, timezone
-from typing import List, Optional
+from typing import Callable, List, Optional
 
 import pytest  # type: ignore
 import responses
@@ -84,6 +85,21 @@ def test_tables(create_table):
     return ids
 
 
+@pytest.fixture(scope="class")
+def supports_test_result_id(client: DataFrameClient) -> bool:
+    """Fixture to determine if the server supports test result IDs."""
+    api_info = client.api_info()
+    result = api_info.operations.create_tables.version >= 2
+    if not result:
+        warnings.warn(
+            (
+                "The selected version of the DataFrame Service does not support test "
+                "result IDs. Tests will not attempt to set a test result ID."
+            )
+        )
+    return result
+
+
 @pytest.mark.enterprise
 @pytest.mark.integration
 class TestDataFrame:
@@ -126,6 +142,43 @@ def test__create_table__metadata_is_correct(
                 properties={},
             ),
         ]
+        assert table_metadata.test_result_id is None
+
+    def test__create_table__supports_test_result(
+        self,
+        client: DataFrameClient,
+        create_table: Callable[[CreateTableRequest], str],
+        supports_test_result_id: bool,
+    ):
+        test_result_id = "Test result" if supports_test_result_id else None
+        id = create_table(
+            CreateTableRequest(
+                columns=[
+                    Column(
+                        name="index",
+                        data_type=DataType.Int32,
+                        column_type=ColumnType.Index,
+                    )
+                ],
+                name="Test table with test result ID",
+                test_result_id=test_result_id,
+            )
+        )
+
+        table = client.get_table_metadata(id)
+        assert table.test_result_id == test_result_id
+
+        page = client.list_tables(take=1, id=[id])
+        assert len(page.tables) == 1
+        assert page.tables[0].test_result_id == test_result_id
+        assert page.continuation_token is None
+
+        page = client.query_tables(
+            QueryTablesRequest(filter="id == @0", substitutions=[id], take=1)
+        )
+        assert len(page.tables) == 1
+        assert page.tables[0].test_result_id == test_result_id
+        assert page.continuation_token is None
 
     def test__get_table__correct_timestamp(self, client: DataFrameClient, create_table):
         id = create_table(basic_table_model)
@@ -188,24 +241,33 @@ def test__query_tables__returns(
         assert len(second_page.tables) == 1
         assert second_page.continuation_token is None
 
-    def test__modify_table__returns(self, client: DataFrameClient, create_table):
+    def test__modify_table__returns(
+        self, client: DataFrameClient, create_table, supports_test_result_id: bool
+    ):
         id = create_table(basic_table_model)
 
         client.modify_table(
             id,
-            ModifyTableRequest(
-                metadata_revision=2,
-                name="Modified table",
-                properties={"cow": "moo"},
-                columns=[
-                    ColumnMetadataPatch(name="index", properties={"sheep": "baa"})
-                ],
+            self._remove_test_result_id_if_not_supported(
+                ModifyTableRequest(
+                    metadata_revision=2,
+                    name="Modified table",
+                    test_result_id="Test result",
+                    properties={"cow": "moo"},
+                    columns=[
+                        ColumnMetadataPatch(name="index", properties={"sheep": "baa"})
+                    ],
+                ),
+                supports_test_result_id,
             ),
         )
         table = client.get_table_metadata(id)
 
         assert table.metadata_revision == 2
         assert table.name == "Modified table"
+        assert table.test_result_id == (
+            "Test result" if supports_test_result_id else None
+        )
         assert table.properties == {"cow": "moo"}
         assert table.columns[0].properties == {"sheep": "baa"}
 
@@ -214,23 +276,41 @@ def test__modify_table__returns(self, client: DataFrameClient, create_table):
 
         assert table.properties == {"cow": "moo", "bee": "buzz"}
         assert table.name == "Modified table"
+        assert table.test_result_id == (
+            "Test result" if supports_test_result_id else None
+        )
 
         client.modify_table(
             id,
-            ModifyTableRequest(
-                metadata_revision=4,
-                name=None,
-                properties={"cow": None},
-                columns=[ColumnMetadataPatch(name="index", properties={"sheep": None})],
+            self._remove_test_result_id_if_not_supported(
+                ModifyTableRequest(
+                    metadata_revision=4,
+                    name=None,
+                    test_result_id=None,
+                    properties={"cow": None},
+                    columns=[
+                        ColumnMetadataPatch(name="index", properties={"sheep": None})
+                    ],
+                ),
+                supports_test_result_id,
             ),
         )
         table = client.get_table_metadata(id)
 
         assert table.metadata_revision == 4
         assert table.name == id
+        assert table.test_result_id is None
         assert table.properties == {"bee": "buzz"}
         assert table.columns[0].properties == {}
 
+    @staticmethod
+    def _remove_test_result_id_if_not_supported(
+        model: ModifyTableRequest, supports_test_result_id: bool
+    ) -> ModifyTableRequest:
+        if not supports_test_result_id:
+            del model.test_result_id
+        return model
+
     def test__delete_table__deletes(self, client: DataFrameClient):
         id = client.create_table(
             basic_table_model
@@ -259,13 +339,16 @@ def test__delete_tables__returns_partial_success(self, client: DataFrameClient):
         assert len(response.error.inner_errors) == 1
 
     def test__modify_tables__modifies_tables(
-        self, client: DataFrameClient, create_table
+        self, client: DataFrameClient, create_table, supports_test_result_id: bool
     ):
         ids = [create_table(basic_table_model) for _ in range(3)]
 
         updates = [
             TableMetadataModification(
-                id=id, name="Modified table", properties={"duck": "quack"}
+                id=id,
+                name="Modified table",
+                test_result_id="Test result" if supports_test_result_id else None,
+                properties={"duck": "quack"},
             )
             for id in ids
         ]
@@ -274,6 +357,9 @@ def test__modify_tables__modifies_tables(
 
         for table in client.list_tables(id=ids).tables:
             assert table.name == "Modified table"
+            assert table.test_result_id == (
+                "Test result" if supports_test_result_id else None
+            )
             assert table.properties == {"duck": "quack"}
 
         updates = [
@@ -286,15 +372,34 @@ def test__modify_tables__modifies_tables(
         )
 
         for table in client.list_tables(id=ids).tables:
+            assert table.name == "Modified table"
+            assert table.test_result_id == (
+                "Test result" if supports_test_result_id else None
+            )
             assert table.properties == {"pig": "oink"}
 
+        if supports_test_result_id:
+            updates = [
+                TableMetadataModification(id=id, test_result_id="") for id in ids
+            ]
+
+            assert client.modify_tables(ModifyTablesRequest(tables=updates)) is None
+
+            for table in client.list_tables(id=ids).tables:
+                assert table.name == "Modified table"
+                assert table.test_result_id is None
+
     def test__modify_tables__returns_partial_success(
-        self, client: DataFrameClient, create_table
+        self, client: DataFrameClient, create_table, supports_test_result_id: bool
     ):
         id = create_table(basic_table_model)
 
         updates = [
-            TableMetadataModification(id=id, name="Modified table")
+            TableMetadataModification(
+                id=id,
+                name="Modified table",
+                test_result_id="Test result" if supports_test_result_id else None,
+            )
             for id in [id, "invalid_id"]
         ]