def test_batch_molecule_graph_data(self):
    (num_nodes_list, num_edge_list) = ([3, 4, 5], [2, 4, 5])
    (num_node_features, num_edge_features) = (32, 32)
    edge_index_list = [np.array([[0, 1], [1, 2]]), np.array([[0, 1, 2, 3], [1, 2, 0, 2]]), np.array([[0, 1, 2, 3, 4], [1, 2, 3, 4, 5]])]
    targets = np.random.random_sample(5)
    molecule_graphs = [MoleculeGraphData(node_features=np.random.random_sample((num_nodes_list[i], num_node_features)), edge_index=edge_index_list[i], targets=targets, edge_features=np.random.random_sample((num_edge_list[i], num_edge_features)), graph_features=None) for i in range(len(num_edge_list))]
    batch = BatchMoleculeGraphData(molecule_graphs)
    assert (batch.num_nodes == sum(num_nodes_list))
    assert (batch.num_node_features == num_node_features)
    assert (batch.num_edges == sum(num_edge_list))
    assert (batch.num_edge_features == num_edge_features)
    assert (batch.targets.shape == (3, 5))
    assert (batch.graph_index.shape == (sum(num_nodes_list),))
def test_batch_molecule_graph_data(self):
    (num_nodes_list, num_edge_list) = ([3, 4, 5], [2, 4, 5])
    (num_node_features, num_edge_features) = (32, 32)
    edge_index_list = [np.array([[0, 1], [1, 2]]), np.array([[0, 1, 2, 3], [1, 2, 0, 2]]), np.array([[0, 1, 2, 3, 4], [1, 2, 3, 4, 5]])]
    targets = np.random.random_sample(5)
    molecule_graphs = [MoleculeGraphData(node_features=np.random.random_sample((num_nodes_list[i], num_node_features)), edge_index=edge_index_list[i], targets=targets, edge_features=np.random.random_sample((num_edge_list[i], num_edge_features)), graph_features=None) for i in range(len(num_edge_list))]
    batch = BatchMoleculeGraphData(molecule_graphs)
    assert (batch.num_nodes == sum(num_nodes_list))
    assert (batch.num_node_features == num_node_features)
    assert (batch.num_edges == sum(num_edge_list))
    assert (batch.num_edge_features == num_edge_features)
    assert (batch.targets.shape == (3, 5))
    assert (batch.graph_index.shape == (sum(num_nodes_list),))