Skip to content

Commit

Permalink
test: Add dataframe-from-existing test
Browse files Browse the repository at this point in the history
  • Loading branch information
j6k4m8 committed Dec 11, 2024
1 parent 213daa9 commit 2062f1a
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 3 deletions.
8 changes: 5 additions & 3 deletions grand/backends/_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,12 @@ def __init__(
use as the node ID
"""
self._directed = directed
self._edge_df = edge_df or pd.DataFrame(
columns=[edge_df_source_column, edge_df_target_column]
self._edge_df = (
edge_df
if edge_df is not None
else pd.DataFrame(columns=[edge_df_source_column, edge_df_target_column])
)
self._node_df = node_df or None
self._node_df = node_df if node_df is not None else None
self._edge_df_source_column = edge_df_source_column
self._edge_df_target_column = edge_df_target_column
self._node_df_id_column = node_df_id_column
Expand Down
34 changes: 34 additions & 0 deletions grand/backends/test_backends.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import pytest
import os
import pandas as pd

import networkx as nx

Expand Down Expand Up @@ -448,3 +449,36 @@ def test_get_density_performance(backend):
for i in range(1000 - 1):
G.nx.add_edge(i, i + 1)
assert nx.density(G.nx) <= 0.005


class TestDataFrameBackend:

def test_can_create_empty(self):
b = DataFrameBackend()
assert b.get_edge_count() == 0
assert b.get_node_count() == 0

b.add_edge("A", "B", {})
assert b.get_edge_count() == 1
assert b.get_node_count() == 2

def test_can_create_from_int_dataframes(self):
# Create an edges DataFrame
edges = pd.DataFrame(
{
"source": [0, 1, 2, 3, 4],
"target": [1, 2, 3, 4, 0],
"weight": [1, 2, 3, 4, 5],
}
)

nodes = pd.DataFrame(
{
"name": [0, 1, 2, 3, 4],
"value": [1, 2, 3, 4, 5],
}
)

b = DataFrameBackend(edge_df=edges, node_df=nodes)
assert b.get_edge_count() == 5
assert b.get_node_count() == 5

0 comments on commit 2062f1a

Please sign in to comment.