Improve knn tests

Author: FlorentinD

graphdatascience/tests/integrationV2/procedure_surface/arrow/similarity/test_knn_arrow_endpoints.py

@@ -53,13 +53,13 @@ def test_knn_stats(knn_endpoints: KnnArrowEndpoints, sample_graph: GraphV2) -> N
     """Test KNN stats operation."""
     result = knn_endpoints.stats(G=sample_graph, node_properties=["prop"], top_k=2)
 
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.compute_millis >= 0
     assert result.pre_processing_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.nodes_compared > 0
-    assert result.similarity_pairs >= 0
+    assert result.similarity_pairs > 0
     assert result.node_pairs_considered > 0
     assert "p50" in result.similarity_distribution
 
@@ -86,14 +86,14 @@ def test_knn_mutate(knn_endpoints: KnnArrowEndpoints, sample_graph: GraphV2) ->
         top_k=2,
     )
 
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.pre_processing_millis >= 0
     assert result.compute_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.mutate_millis >= 0
     assert result.relationships_written == sample_graph.node_count() * 2
-    assert result.node_pairs_considered >= 0
+    assert result.node_pairs_considered > 0
 
 
 @pytest.mark.db_integration
@@ -108,14 +108,14 @@ def test_knn_write(arrow_client: AuthenticatedArrowClient, query_runner: QueryRu
     )
 
     assert isinstance(result, KnnWriteResult)
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.pre_processing_millis >= 0
     assert result.compute_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.write_millis >= 0
     assert result.relationships_written == db_graph.node_count() * 2
-    assert result.node_pairs_considered >= 0
+    assert result.node_pairs_considered > 0
 
     # Check that relationships were written to the database
     count_result = query_runner.run_cypher("MATCH ()-[r:SIMILAR]->() RETURN COUNT(r) AS count")

graphdatascience/tests/integrationV2/procedure_surface/arrow/similarity/test_knn_filtered_arrow_endpoints.py

@@ -63,11 +63,11 @@ def test_knn_filtered_stats(knn_filtered_endpoints: KnnFilteredArrowEndpoints, s
     assert result.post_processing_millis >= 0
     assert result.nodes_compared >= 0
     assert result.similarity_pairs >= 0
-    assert result.similarity_distribution is not None
-    assert isinstance(result.did_converge, bool)
+    assert "p50" in result.similarity_distribution
+    assert result.did_converge
     assert result.ran_iterations >= 0
     assert result.node_pairs_considered >= 0
-    assert result.configuration is not None
+    assert "concurrency" in result.configuration
 
 
 def test_knn_filtered_stream(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
@@ -99,12 +99,12 @@ def test_knn_filtered_mutate(knn_filtered_endpoints: KnnFilteredArrowEndpoints,
     assert result.mutate_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.nodes_compared >= 0
-    assert result.relationships_written >= 0
-    assert result.similarity_distribution is not None
-    assert isinstance(result.did_converge, bool)
-    assert result.ran_iterations >= 0
+    assert result.relationships_written > 0
+    assert "p50" in result.similarity_distribution
+    assert result.did_converge
+    assert result.ran_iterations > 0
     assert result.node_pairs_considered >= 0
-    assert result.configuration is not None
+    assert "concurrency" in result.configuration
 
 
 def test_knn_filtered_write(
@@ -124,18 +124,17 @@ def test_knn_filtered_write(
         target_node_filter="TargetNode",
     )
 
-    assert isinstance(result, KnnWriteResult)
     assert result.pre_processing_millis >= 0
     assert result.compute_millis >= 0
     assert result.write_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.nodes_compared >= 0
-    assert result.relationships_written >= 0
-    assert isinstance(result.did_converge, bool)
-    assert result.ran_iterations >= 0
+    assert result.relationships_written > 0
+    assert result.did_converge
+    assert result.ran_iterations > 0
     assert result.node_pairs_considered >= 0
-    assert result.similarity_distribution is not None
-    assert result.configuration is not None
+    assert "p50" in result.similarity_distribution
+    assert "concurrency" in result.configuration
 
 
 def test_knn_filtered_estimate(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:

graphdatascience/tests/integrationV2/procedure_surface/cypher/similarity/test_knn_cypher_endpoints.py

@@ -39,37 +39,31 @@ def knn_endpoints(query_runner: QueryRunner) -> Generator[KnnCypherEndpoints, No
 
 
 def test_knn_stats(knn_endpoints: KnnCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN stats operation."""
     result = knn_endpoints.stats(G=sample_graph, node_properties=["prop"], top_k=2)
 
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.compute_millis > 0
     assert result.pre_processing_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.nodes_compared > 0
     assert result.similarity_pairs == 8
     assert result.node_pairs_considered > 0
-    assert "p50" in result.similarity_distribution or "p10" in result.similarity_distribution
+    assert "p50" in result.similarity_distribution
 
 
 def test_knn_stream(knn_endpoints: KnnCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN stream operation."""
     result_df = knn_endpoints.stream(
         G=sample_graph,
         node_properties=["prop"],
         top_k=2,
     )
 
-    assert "node1" in result_df.columns
-    assert "node2" in result_df.columns
-    assert "similarity" in result_df.columns
-    assert len(result_df.columns) == 3
+    assert set(result_df.columns) == {"node1", "node2", "similarity"}
     assert len(result_df) == 8
 
 
 def test_knn_mutate(knn_endpoints: KnnCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN mutate operation."""
     result = knn_endpoints.mutate(
         G=sample_graph,
         mutate_relationship_type="SIMILAR",
@@ -78,14 +72,14 @@ def test_knn_mutate(knn_endpoints: KnnCypherEndpoints, sample_graph: GraphV2) ->
         top_k=2,
     )
 
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.pre_processing_millis >= 0
     assert result.compute_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.mutate_millis >= 0
     assert result.relationships_written == 8
-    assert result.node_pairs_considered >= 0
+    assert result.node_pairs_considered > 0
 
 
 def test_knn_estimate(knn_endpoints: KnnCypherEndpoints, sample_graph: GraphV2) -> None:

graphdatascience/tests/integrationV2/procedure_surface/cypher/similarity/test_knn_filtered_cypher_endpoints.py

@@ -41,7 +41,6 @@ def knn_filtered_endpoints(query_runner: QueryRunner) -> Generator[KnnFilteredCy
 
 
 def test_knn_filtered_stats(knn_filtered_endpoints: KnnFilteredCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN filtered stats operation."""
     result = knn_filtered_endpoints.stats(
         G=sample_graph,
         node_properties=["prop"],
@@ -50,20 +49,19 @@ def test_knn_filtered_stats(knn_filtered_endpoints: KnnFilteredCypherEndpoints,
         target_node_filter="TargetNode",
     )
 
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.compute_millis > 0
     assert result.pre_processing_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.nodes_compared > 0
-    assert result.similarity_pairs >= 0
-    assert result.similarity_distribution is not None
-    assert result.node_pairs_considered >= 0
-    assert result.configuration is not None
+    assert result.similarity_pairs > 0
+    assert "p50" in result.similarity_distribution
+    assert result.node_pairs_considered > 0
+    assert "concurrency" in result.configuration
 
 
 def test_knn_filtered_stream(knn_filtered_endpoints: KnnFilteredCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN filtered stream operation."""
     result = knn_filtered_endpoints.stream(
         G=sample_graph,
         node_properties=["prop"],
@@ -72,15 +70,11 @@ def test_knn_filtered_stream(knn_filtered_endpoints: KnnFilteredCypherEndpoints,
         target_node_filter="TargetNode",
     )
 
-    assert len(result) >= 0
-    if len(result) > 0:
-        assert "node1" in result.columns
-        assert "node2" in result.columns
-        assert "similarity" in result.columns
+    assert set(result.columns) == {"node1", "node2", "similarity"}
+    assert len(result) >= 4
 
 
 def test_knn_filtered_mutate(knn_filtered_endpoints: KnnFilteredCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN filtered mutate operation."""
     result = knn_filtered_endpoints.mutate(
         G=sample_graph,
         node_properties=["prop"],
@@ -91,21 +85,20 @@ def test_knn_filtered_mutate(knn_filtered_endpoints: KnnFilteredCypherEndpoints,
         target_node_filter="TargetNode",
     )
 
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.compute_millis > 0
     assert result.mutate_millis >= 0
     assert result.pre_processing_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.nodes_compared > 0
-    assert result.relationships_written >= 0
-    assert result.similarity_distribution is not None
-    assert result.node_pairs_considered >= 0
-    assert result.configuration is not None
+    assert result.relationships_written > 0
+    assert "p50" in result.similarity_distribution
+    assert result.node_pairs_considered > 0
+    assert "concurrency" in result.configuration
 
 
 def test_knn_filtered_write(knn_filtered_endpoints: KnnFilteredCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN filtered write operation."""
     result = knn_filtered_endpoints.write(
         G=sample_graph,
         node_properties=["prop"],
@@ -116,21 +109,20 @@ def test_knn_filtered_write(knn_filtered_endpoints: KnnFilteredCypherEndpoints,
         target_node_filter="TargetNode",
     )
 
-    assert result.ran_iterations >= 0
-    assert result.did_converge in [True, False]
+    assert result.ran_iterations > 0
+    assert result.did_converge
     assert result.compute_millis > 0
     assert result.write_millis >= 0
     assert result.pre_processing_millis >= 0
     assert result.post_processing_millis >= 0
     assert result.nodes_compared > 0
-    assert result.relationships_written >= 0
-    assert result.similarity_distribution is not None
-    assert result.node_pairs_considered >= 0
-    assert result.configuration is not None
+    assert result.relationships_written > 0
+    assert "p50" in result.similarity_distribution
+    assert result.node_pairs_considered > 0
+    assert "concurrency" in result.configuration
 
 
 def test_knn_filtered_estimate(knn_filtered_endpoints: KnnFilteredCypherEndpoints, sample_graph: GraphV2) -> None:
-    """Test KNN filtered estimation operation."""
     result = knn_filtered_endpoints.estimate(
         G=sample_graph,
         node_properties=["prop"],