Fix arrow v2 unit testing

* remove usage of unittest
* avoid skipping too many tests

Author: FlorentinD

graphdatascience/tests/integrationV2/conftest.py

@@ -9,7 +9,9 @@ def pytest_collection_modifyitems(config: Any, items: Any) -> None:
     if not config.getoption("--include-integration-v2"):
         skip_v2 = pytest.mark.skip(reason="need --include-integration-v2 option to run")
         for item in items:
-            item.add_marker(skip_v2)
+            # otherwise would also skip lots of other test
+            if "integrationV2" in str(item.fspath):
+                item.add_marker(skip_v2)
 
 
 # best used with pytest --basetemp=tmp/pytest for easy access to logs

graphdatascience/tests/unit/arrow_client/V2/test_job_client.py

@@ -1,109 +1,112 @@
 import json
-import unittest
-from unittest.mock import MagicMock
+
+from pytest_mock import MockerFixture
 
 from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
 from graphdatascience.arrow_client.v2.job_client import JobClient
 from graphdatascience.tests.unit.arrow_client.arrow_test_utils import ArrowTestResult
 
 
-class TestJobClient(unittest.TestCase):
-    def setUp(self) -> None:
-        self.mock_client = MagicMock()
+def test_run_job(mocker: MockerFixture) -> None:
+    mock_client = mocker.Mock()
+    job_id = "test-job-123"
+    endpoint = "v2/test.endpoint"
+    config = {"param1": "value1", "param2": 42}
+
+    mock_client.do_action_with_retry.return_value = iter([ArrowTestResult({"jobId": job_id})])
+
+    result = JobClient.run_job(mock_client, endpoint, config)
+
+    expected_config = json.dumps(config).encode("utf-8")
+    mock_client.do_action_with_retry.assert_called_once_with(endpoint, expected_config)
+    assert result == job_id
 
-    def test_run_job(self) -> None:
-        job_id = "test-job-123"
-        endpoint = "v2/test.endpoint"
-        config = {"param1": "value1", "param2": 42}
 
-        self.mock_client.do_action_with_retry.return_value = iter([ArrowTestResult({"jobId": job_id})])
+def test_run_job_and_wait(mocker: MockerFixture) -> None:
+    mock_client = mocker.Mock()
+    job_id = "test-job-456"
+    endpoint = "v2/test.endpoint"
+    config = {"param": "value"}
 
-        result = JobClient.run_job(self.mock_client, endpoint, config)
+    job_id_config = JobIdConfig(jobId=job_id)
 
-        expected_config = json.dumps(config).encode("utf-8")
-        self.mock_client.do_action_with_retry.assert_called_once_with(endpoint, expected_config)
-        self.assertEqual(result, job_id)
+    status = JobStatus(
+        jobId=job_id,
+        progress=1.0,
+        status="Done",
+    )
 
-    def test_run_job_and_wait(
-        self,
-    ) -> None:
-        job_id = "test-job-456"
-        endpoint = "v2/test.endpoint"
-        config = {"param": "value"}
+    do_action_with_retry = mocker.Mock()
+    do_action_with_retry.side_effect = [
+        iter([ArrowTestResult(job_id_config.dump_camel())]),
+        iter([ArrowTestResult(status.dump_camel())]),
+    ]
 
-        job_id_config = JobIdConfig(jobId=job_id)
+    mock_client.do_action_with_retry = do_action_with_retry
 
-        status = JobStatus(
-            jobId=job_id,
-            progress=1.0,
-            status="Done",
-        )
+    result = JobClient.run_job_and_wait(mock_client, endpoint, config)
 
-        do_action_with_retry = MagicMock()
-        do_action_with_retry.side_effect = [
-            iter([ArrowTestResult(job_id_config.dump_camel())]),
-            iter([ArrowTestResult(status.dump_camel())]),
-        ]
+    do_action_with_retry.assert_called_with("v2/jobs.status", job_id_config.dump_json().encode("utf-8"))
+    assert result == job_id
 
-        self.mock_client.do_action_with_retry = do_action_with_retry
 
-        result = JobClient.run_job_and_wait(self.mock_client, endpoint, config)
+def test_wait_for_job_completes_immediately(mocker: MockerFixture) -> None:
+    mock_client = mocker.Mock()
+    job_id = "test-job-789"
 
-        do_action_with_retry.assert_called_with("v2/jobs.status", job_id_config.dump_json().encode("utf-8"))
-        self.assertEqual(result, job_id)
+    status = JobStatus(
+        jobId=job_id,
+        progress=1.0,
+        status="Done",
+    )
 
-    def test_wait_for_job_completes_immediately(self) -> None:
-        job_id = "test-job-789"
+    mock_client.do_action_with_retry.return_value = iter([ArrowTestResult(status.dump_camel())])
 
-        status = JobStatus(
-            jobId=job_id,
-            progress=1.0,
-            status="Done",
-        )
+    JobClient.wait_for_job(mock_client, job_id)
 
-        self.mock_client.do_action_with_retry.return_value = iter([ArrowTestResult(status.dump_camel())])
+    mock_client.do_action_with_retry.assert_called_once_with(
+        "v2/jobs.status", JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+    )
 
-        JobClient.wait_for_job(self.mock_client, job_id)
 
-        self.mock_client.do_action_with_retry.assert_called_once_with(
-            "v2/jobs.status", JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
-        )
+def test_wait_for_job_waits_for_completion(mocker: MockerFixture) -> None:
+    mock_client = mocker.Mock()
+    job_id = "test-job-waiting"
+    status_running = JobStatus(
+        jobId=job_id,
+        progress=0.5,
+        status="RUNNING",
+    )
+    status_done = JobStatus(
+        jobId=job_id,
+        progress=1.0,
+        status="Done",
+    )
 
-    def test_wait_for_job_waits_for_completion(self) -> None:
-        job_id = "test-job-waiting"
-        status_running = JobStatus(
-            jobId=job_id,
-            progress=0.5,
-            status="RUNNING",
-        )
-        status_done = JobStatus(
-            jobId=job_id,
-            progress=1.0,
-            status="Done",
-        )
+    do_action_with_retry = mocker.Mock()
+    do_action_with_retry.side_effect = [
+        iter([ArrowTestResult(status_running.dump_camel())]),
+        iter([ArrowTestResult(status_done.dump_camel())]),
+    ]
 
-        do_action_with_retry = MagicMock()
-        do_action_with_retry.side_effect = [
-            iter([ArrowTestResult(status_running.dump_camel())]),
-            iter([ArrowTestResult(status_done.dump_camel())]),
-        ]
+    mock_client.do_action_with_retry = do_action_with_retry
 
-        self.mock_client.do_action_with_retry = do_action_with_retry
+    JobClient.wait_for_job(mock_client, job_id)
 
-        JobClient.wait_for_job(self.mock_client, job_id)
+    assert mock_client.do_action_with_retry.call_count == 2
 
-        self.assertEqual(self.mock_client.do_action_with_retry.call_count, 2)
 
-    def test_get_summary(self) -> None:
-        # Setup
-        job_id = "summary-job-123"
-        expected_summary = {"nodeCount": 100, "relationshipCount": 200, "requiredMemory": "1GB"}
+def test_get_summary(mocker: MockerFixture) -> None:
+    mock_client = mocker.Mock()
+    # Setup
+    job_id = "summary-job-123"
+    expected_summary = {"nodeCount": 100, "relationshipCount": 200, "requiredMemory": "1GB"}
 
-        self.mock_client.do_action_with_retry.return_value = iter([ArrowTestResult(expected_summary)])
+    mock_client.do_action_with_retry.return_value = iter([ArrowTestResult(expected_summary)])
 
-        result = JobClient.get_summary(self.mock_client, job_id)
+    result = JobClient.get_summary(mock_client, job_id)
 
-        self.mock_client.do_action_with_retry.assert_called_once_with(
-            "v2/results.summary", JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
-        )
-        self.assertEqual(result, expected_summary)
+    mock_client.do_action_with_retry.assert_called_once_with(
+        "v2/results.summary", JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+    )
+    assert result == expected_summary

graphdatascience/tests/unit/arrow_client/V2/test_mutation_client.py

@@ -1,27 +1,25 @@
-import unittest
-from unittest.mock import MagicMock
+import json
+
+from pytest_mock import MockerFixture
 
-from graphdatascience.arrow_client.v2.api_types import MutateResult
 from graphdatascience.arrow_client.v2.mutation_client import MutationClient
 from graphdatascience.tests.unit.arrow_client.arrow_test_utils import ArrowTestResult
 
 
-class TestMutationClient(unittest.TestCase):
-    def setUp(self) -> None:
-        self.mock_client = MagicMock()
+def test_mutate_node_property_success(mocker: MockerFixture) -> None:
+    job_id = "test-job-123"
+    arrow_mutation_result = {"nodePropertiesWritten": 42, "relationshipsWritten": 1337}
 
-    def test_mutate_node_property_success(self) -> None:
-        job_id = "test-job-123"
-        expected_mutation_result = MutateResult(nodePropertiesWritten=42, relationshipsWritten=1337)
+    mock_client = mocker.Mock()
+    mock_client.do_action_with_retry.return_value = iter([ArrowTestResult(arrow_mutation_result)])
 
-        self.mock_client.do_action_with_retry.return_value = iter(
-            [ArrowTestResult(expected_mutation_result.dump_camel())]
-        )
+    result = MutationClient.mutate_node_property(mock_client, job_id, "propertyName")
 
-        result = MutationClient.mutate_node_property(self.mock_client, job_id, "propertyName")
+    assert result.node_properties_written == 42
+    assert result.relationships_written == 1337
+    assert result.mutate_millis > 0
 
-        assert result == expected_mutation_result
+    args, _ = mock_client.do_action_with_retry.call_args
 
-        self.mock_client.do_action_with_retry.assert_called_once_with(
-            MutationClient.MUTATE_ENDPOINT, b'{"jobId": "test-job-123", "mutateProperty": "propertyName"}'
-        )
+    assert args[0] == MutationClient.MUTATE_ENDPOINT
+    assert json.loads(args[1]) == {"jobId": "test-job-123", "mutateProperty": "propertyName"}

graphdatascience/tests/unit/arrow_client/V2/test_write_back_client.py

@@ -1,18 +1,18 @@
 from typing import Optional
-from unittest.mock import Mock
 
 import pytest
 from pandas import DataFrame
+from pytest_mock import MockerFixture
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.write_back_client import WriteBackClient
 from graphdatascience.tests.unit.conftest import DEFAULT_SERVER_VERSION, CollectingQueryRunner
 
 
 @pytest.fixture
-def mock_arrow_client() -> AuthenticatedArrowClient:
-    client = Mock(spec=AuthenticatedArrowClient)
-    client.connection_info.return_value = Mock(host="localhost", port=8080, encrypted=False)
+def mock_arrow_client(mocker: MockerFixture) -> AuthenticatedArrowClient:
+    client = mocker.Mock(spec=AuthenticatedArrowClient)
+    client.connection_info.return_value = mocker.Mock(host="localhost", port=8080, encrypted=False)
     client.request_token.return_value = "test_token"
     return client
 
@@ -44,12 +44,12 @@ def test_arrow_configuration(write_back_client: WriteBackClient, mock_arrow_clie
     assert config == expected_config
 
 
-def test_write_calls_run_write_back(write_back_client: WriteBackClient) -> None:
+def test_write_calls_run_write_back(write_back_client: WriteBackClient, mocker: MockerFixture) -> None:
     graph_name = "test_graph"
     job_id = "123"
     concurrency: Optional[int] = 4
 
-    write_back_client._write_protocol.run_write_back = Mock()  # type: ignore
+    write_back_client._write_protocol.run_write_back = mocker.Mock()  # type: ignore
 
     duration = write_back_client.write(graph_name, job_id, concurrency)
 

tox.ini

@@ -114,4 +114,3 @@ passenv =
 	BUILD_NUMBER
 commands =
 	pytest graphdatascience/tests/integrationV2 --include-integration-v2
-	pytest graphdatascience/tests/unit/procedure_surface