def test_OnDiskDataset_TVTSet_HeteroItemSet_id_label():
    """Test TVTSet which returns HeteroItemSet with IDs and labels."""
    with tempfile.TemporaryDirectory() as test_dir:
        train_ids = np.arange(1000)
        train_labels = np.random.randint(0, 10, size=1000)
        train_data = np.vstack([train_ids, train_labels]).T
        train_path = os.path.join(test_dir, "train.npy")
        np.save(train_path, train_data)

        validation_ids = np.arange(1000, 2000)
        validation_labels = np.random.randint(0, 10, size=1000)
        validation_data = np.vstack([validation_ids, validation_labels]).T
        validation_path = os.path.join(test_dir, "validation.npy")
        np.save(validation_path, validation_data)

        test_ids = np.arange(2000, 3000)
        test_labels = np.random.randint(0, 10, size=1000)
        test_data = np.vstack([test_ids, test_labels]).T
        test_path = os.path.join(test_dir, "test.npy")
        np.save(test_path, test_data)

        yaml_content = f"""
            tasks:
              - name: node_classification
                train_set:
                  - type: paper
                    data:
                      - name: seeds
                        format: numpy
                        in_memory: true
                        path: {train_path}
                  - type: author
                    data:
                      - name: seeds
                        format: numpy
                        path: {train_path}
                validation_set:
                  - type: paper
                    data:
                      - name: seeds
                        format: numpy
                        path: {validation_path}
                  - type: author
                    data:
                      - name: seeds
                        format: numpy
                        path: {validation_path}
                test_set:
                  - type: paper
                    data:
                      - name: seeds
                        format: numpy
                        in_memory: false
                        path: {test_path}
                  - type: author
                    data:
                      - name: seeds
                        format: numpy
                        path: {test_path}
        """
        dataset = write_yaml_and_load_dataset(yaml_content, test_dir)

        # Verify train set.
        train_set = dataset.tasks[0].train_set
        assert len(train_set) == 2000
        assert isinstance(train_set, gb.HeteroItemSet)
        for i, item in enumerate(train_set):
            assert isinstance(item, dict)
            assert len(item) == 1
            key = list(item.keys())[0]
            assert key in ["paper", "author"]
            id, label = item[key]
            assert id == train_ids[i % 1000]
            assert label == train_labels[i % 1000]
        assert train_set.names == ("seeds",)
        train_set = None

        # Verify validation set.
        validation_set = dataset.tasks[0].validation_set
        assert len(validation_set) == 2000
        assert isinstance(validation_set, gb.HeteroItemSet)
        for i, item in enumerate(validation_set):
            assert isinstance(item, dict)
            assert len(item) == 1
            key = list(item.keys())[0]
            assert key in ["paper", "author"]
            id, label = item[key]
            assert id == validation_ids[i % 1000]
            assert label == validation_labels[i % 1000]
        assert validation_set.names == ("seeds",)
        validation_set = None

        # Verify test set.
        test_set = dataset.tasks[0].test_set
        assert len(test_set) == 2000
        assert isinstance(test_set, gb.HeteroItemSet)
        for i, item in enumerate(test_set):
            assert isinstance(item, dict)
            assert len(item) == 1
            key = list(item.keys())[0]
            assert key in ["paper", "author"]
            id, label = item[key]
            assert id == test_ids[i % 1000]
            assert label == test_labels[i % 1000]
        assert test_set.names == ("seeds",)
        test_set = None
        dataset = None
Test TVTSet which returns HeteroItemSet with IDs and labels.
tests/python/pytorch/graphbolt/impl/test_ondisk_dataset.py