def test_OnDiskDataset_all_nodes_set_hetero():
    """Test heterograph's all nodes set of OnDiskDataset."""
    (
        csc_indptr,
        indices,
        node_type_offset,
        type_per_edge,
        node_type_to_id,
        edge_type_to_id,
    ) = gbt.random_hetero_graph(1000, 10 * 1000, 3, 4)
    graph = gb.fused_csc_sampling_graph(
        csc_indptr,
        indices,
        node_type_offset=node_type_offset,
        type_per_edge=type_per_edge,
        node_type_to_id=node_type_to_id,
        edge_type_to_id=edge_type_to_id,
        edge_attributes=None,
    )

    with tempfile.TemporaryDirectory() as test_dir:
        graph_path = os.path.join(test_dir, "fused_csc_sampling_graph.pt")
        torch.save(graph, graph_path)

        yaml_content = f"""
            graph_topology:
              type: FusedCSCSamplingGraph
              path: {graph_path}
        """
        dataset = write_yaml_and_load_dataset(yaml_content, test_dir)
        all_nodes_set = dataset.all_nodes_set
        assert isinstance(all_nodes_set, gb.HeteroItemSet)
        assert all_nodes_set.names == ("seeds",)
        for i, item in enumerate(all_nodes_set):
            assert len(item) == 1
            assert isinstance(item, dict)

        dataset = None
Test heterograph's all nodes set of OnDiskDataset.
tests/python/pytorch/graphbolt/impl/test_ondisk_dataset.py