From 7b27580a44ee216bbf4af2e65f033d1b0fc7c03d Mon Sep 17 00:00:00 2001
From: chrisiacovella <chris.iacovella@gmail.com>
Date: Thu, 14 Sep 2023 23:06:48 -0700
Subject: [PATCH] Working on updating unit tests.

---
 modelforge/dataset/qm9.py        | 12 +++++++-----
 modelforge/tests/test_dataset.py | 10 +++++-----
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/modelforge/dataset/qm9.py b/modelforge/dataset/qm9.py
index 027e74fb..58d3bd1c 100644
--- a/modelforge/dataset/qm9.py
+++ b/modelforge/dataset/qm9.py
@@ -31,13 +31,13 @@ class QM9Dataset(HDF5Dataset):
     _property_names = PropertyNames(
         "atomic_numbers",
         "geometry",
-        "return_energy",
+        "internal_energy_at_0K",
     )
 
     _available_properties = [
         "geometry",
         "atomic_numbers",
-        "return_energy",
+        "internal_energy_at_0K",
     ]  # NOTE: Any way to set this automatically?
 
     def __init__(
@@ -68,7 +68,7 @@ def __init__(
         _default_properties_of_interest = [
             "geometry",
             "atomic_numbers",
-            "return_energy",
+            "internal_energy_at_0K",
         ]  # NOTE: Default values
 
         self._properties_of_interest = _default_properties_of_interest
@@ -82,8 +82,10 @@ def __init__(
         )
         self.dataset_name = dataset_name
         self.for_unit_testing = for_unit_testing
-        self.test_id = "17oZ07UOxv2fkEmu-d5mLk6aGIuhV0mJ7"
-        self.full_id = "1_bSdQjEvI67Tk_LKYbW0j8nmggnb5MoU"
+        # self.test_id = "17oZ07UOxv2fkEmu-d5mLk6aGIuhV0mJ7"
+        # self.full_id = "1_bSdQjEvI67Tk_LKYbW0j8nmggnb5MoU"
+        self.test_id = "18C9Iq_7VZLx0gZbJYje8X6tybZb5m3JY"
+        self.full_id = "1damjPgjKviTogDJ2UJvhYjyBZxGvRPP-"
 
     @property
     def properties_of_interest(self) -> List[str]:
diff --git a/modelforge/tests/test_dataset.py b/modelforge/tests/test_dataset.py
index 17254511..b565984d 100644
--- a/modelforge/tests/test_dataset.py
+++ b/modelforge/tests/test_dataset.py
@@ -64,7 +64,7 @@ def test_different_properties_of_interest(dataset):
     assert data.properties_of_interest == [
         "geometry",
         "atomic_numbers",
-        "return_energy",
+        "internal_energy_at_0K",
     ]
 
     dataset = factory.create_dataset(data)
@@ -72,9 +72,9 @@ def test_different_properties_of_interest(dataset):
     assert isinstance(raw_data_item, dict)
     assert len(raw_data_item) == 4
 
-    data.properties_of_interest = ["return_energy", "geometry"]
+    data.properties_of_interest = ["internal_energy_at_0K", "geometry"]
     assert data.properties_of_interest == [
-        "return_energy",
+        "internal_energy_at_0K",
         "geometry",
     ]
 
@@ -82,7 +82,7 @@ def test_different_properties_of_interest(dataset):
     raw_data_item = dataset[0]
     print(raw_data_item)
     assert isinstance(raw_data_item, dict)
-    assert len(raw_data_item) != 3  
+    assert len(raw_data_item) != 3
 
 
 @pytest.mark.parametrize("dataset", DATASETS)
@@ -165,7 +165,7 @@ def test_dataset_generation(dataset):
         pass
 
     # the dataloader automatically splits and batches the dataset
-    # for the trianing set it batches the 80 datapoints in
+    # for the training set it batches the 80 datapoints in
     # a batch of 64 and a batch of 16 samples
     assert len(train_dataloader) == 2  # nr of batches
     v = [v_ for v_ in train_dataloader]