added test cases

Author: yh-yao

fedgraph/trainer_class.py

@@ -186,6 +186,9 @@ def __init__(
         # self.class_num = class_num
         self.args = args
         self.model = None
+        self.optimizer = None
+        self.global_node_num = None
+        self.class_num = None
         self.feature_aggregation = None
         if self.args.method == "FedAvg":
             # print("Loading feature as the feature aggregation for fedavg method")
@@ -276,6 +279,9 @@ def update_params(self, params: tuple, current_global_epoch: int) -> None:
             The current global epoch number.
         """
         # load global parameter from global server
+        if self.model is None:
+            return
+            
         self.model.to("cpu")
         for (
             p,
@@ -314,17 +320,22 @@ def get_local_feature_sum(self) -> torch.Tensor:
         normalized_sum : torch.Tensor
             The normalized sum of features of 1-hop neighbors for each node
         """
+        # Use global_node_num if available, otherwise infer from communicate_node_index
+        global_node_num = getattr(self, 'global_node_num', None)
+        if global_node_num is None:
+            global_node_num = self.communicate_node_index.max().item() + 1
+            
         # Create a large matrix with known local node features
         new_feature_for_trainer = torch.zeros(
-            self.global_node_num, self.features.shape[1]
+            global_node_num, self.features.shape[1]
         ).to(self.device)
         new_feature_for_trainer[self.local_node_index] = self.features
 
         # Sum of features of all 1-hop nodes for each node
         one_hop_neighbor_feature_sum = get_1hop_feature_sum(
             new_feature_for_trainer, self.adj, self.device
         )
-        if self.args.use_encryption:
+        if hasattr(self.args, 'use_encryption') and self.args.use_encryption:
             print(
                 f"Trainer {self.rank} - Original feature sum (first 10 and last 10 elements): "
                 f"{one_hop_neighbor_feature_sum.flatten()[:10].tolist()} ... {one_hop_neighbor_feature_sum.flatten()[-10:].tolist()}"
@@ -522,6 +533,7 @@ def train(self, current_global_round: int) -> None:
             )
 
         self.feature_aggregation = self.feature_aggregation.to(self.device)
+        data = None
         if hasattr(self.args, "batch_size") and self.args.batch_size > 0:
             # batch preparation
             train_mask = torch.zeros(
@@ -541,6 +553,8 @@ def train(self, current_global_round: int) -> None:
                 train_mask=train_mask,
                 y=node_labels,
             )
+        loss_train = 0.0
+        acc_train = 0.0
         for iteration in range(self.local_step):
             self.model.train()
             if hasattr(self.args, "batch_size") and self.args.batch_size > 0:
@@ -621,13 +635,18 @@ def local_test(self) -> list:
         (list) : list
             A list containing the test loss and accuracy [local_test_loss, local_test_acc].
         """
+        if self.model is None or self.feature_aggregation is None:
+            return [0.0, 0.0]
+            
         local_test_loss, local_test_acc = test(
             self.model,
             self.feature_aggregation,
             self.adj,
             self.test_labels,
             self.idx_test,
         )
+        self.test_losses.append(local_test_loss)
+        self.test_accs.append(local_test_acc)
         return [local_test_loss, local_test_acc]
 
     def get_params(self) -> tuple:
@@ -639,8 +658,11 @@ def get_params(self) -> tuple:
         (tuple) : tuple
             A tuple containing the current parameters of the model.
         """
-        self.optimizer.zero_grad(set_to_none=True)
-        return tuple(self.model.parameters())
+        if self.optimizer is not None:
+            self.optimizer.zero_grad(set_to_none=True)
+        if self.model is not None:
+            return tuple(self.model.parameters())
+        return ()
 
     def get_all_loss_accuray(self) -> list:
         """

pytest.ini

@@ -0,0 +1,16 @@
+[pytest]
+testpaths = tests
+python_files = test_*.py *_test.py
+python_classes = Test*
+python_functions = test_*
+addopts = -v --tb=short --strict-markers
+markers =
+    unit: Unit tests
+    integration: Integration tests
+    slow: Slow running tests
+    gpu: Tests requiring GPU
+    ray: Tests requiring Ray cluster
+filterwarnings =
+    ignore::DeprecationWarning
+    ignore::UserWarning
+    ignore::PendingDeprecationWarning

tests/__init__.py

@@ -0,0 +1,123 @@
+import pytest
+import torch
+import numpy as np
+import tempfile
+import os
+from unittest.mock import Mock, patch, MagicMock
+from torch_geometric.data import Data
+import shutil
+
+@pytest.fixture
+def temp_dir():
+    """Create a temporary directory for test files."""
+    temp_dir = tempfile.mkdtemp()
+    yield temp_dir
+    shutil.rmtree(temp_dir)
+
+@pytest.fixture
+def sample_graph_data():
+    """Create sample graph data for testing."""
+    num_nodes = 10
+    num_features = 5
+    num_classes = 3
+    
+    # Create node features
+    x = torch.randn(num_nodes, num_features)
+    
+    # Create edges (simple chain graph)
+    edge_index = torch.tensor([[i, i+1] for i in range(num_nodes-1)]).t().contiguous()
+    
+    # Create labels
+    y = torch.randint(0, num_classes, (num_nodes,))
+    
+    data = Data(x=x, edge_index=edge_index, y=y)
+    return data
+
+@pytest.fixture
+def mock_args():
+    """Create mock arguments for testing."""
+    args = Mock()
+    args.num_clients = 5
+    args.dataset = 'Cora'
+    args.method = 'FedAvg'
+    args.num_rounds = 10
+    args.local_epochs = 5
+    args.lr = 0.01
+    args.hidden = 64
+    args.dropout = 0.5
+    args.device = 'cpu'
+    args.seed = 42
+    args.data_path = '/tmp/test_data'
+    args.split_type = 'louvain'
+    args.alpha = 0.5
+    args.beta = 1.0
+    args.num_workers = 1
+    args.ray_init_address = None
+    args.ray_dashboard_port = 8265
+    args.he = False
+    args.dp = False
+    return args
+
+@pytest.fixture
+def mock_ray_cluster():
+    """Mock Ray cluster for testing."""
+    with patch('ray.init') as mock_init, \
+         patch('ray.get') as mock_get, \
+         patch('ray.put') as mock_put, \
+         patch('ray.remote') as mock_remote:
+        
+        mock_init.return_value = None
+        mock_get.side_effect = lambda x: x
+        mock_put.side_effect = lambda x: x
+        mock_remote.side_effect = lambda x: x
+        
+        yield {
+            'init': mock_init,
+            'get': mock_get,
+            'put': mock_put,
+            'remote': mock_remote
+        }
+
+@pytest.fixture
+def sample_dataset_splits():
+    """Create sample dataset splits for federated learning."""
+    num_clients = 3
+    num_nodes_per_client = 5
+    
+    splits = {}
+    for i in range(num_clients):
+        client_data = {
+            'train_mask': torch.zeros(num_nodes_per_client, dtype=torch.bool),
+            'val_mask': torch.zeros(num_nodes_per_client, dtype=torch.bool),
+            'test_mask': torch.zeros(num_nodes_per_client, dtype=torch.bool),
+            'node_list': list(range(i * num_nodes_per_client, (i + 1) * num_nodes_per_client))
+        }
+        # Set some nodes for training
+        client_data['train_mask'][:3] = True
+        client_data['val_mask'][3:4] = True
+        client_data['test_mask'][4:5] = True
+        
+        splits[f'client_{i}'] = client_data
+    
+    return splits
+
+@pytest.fixture
+def mock_model():
+    """Create a mock GNN model for testing."""
+    model = Mock()
+    model.parameters.return_value = [torch.randn(10, 5, requires_grad=True)]
+    model.state_dict.return_value = {'layer.weight': torch.randn(10, 5)}
+    model.load_state_dict = Mock()
+    model.train = Mock()
+    model.eval = Mock()
+    return model
+
+@pytest.fixture
+def mock_optimizer():
+    """Create a mock optimizer for testing."""
+    optimizer = Mock()
+    optimizer.zero_grad = Mock()
+    optimizer.step = Mock()
+    optimizer.state_dict.return_value = {}
+    optimizer.load_state_dict = Mock()
+    return optimizer