From 4a8d9706be6dc8bc1bbb43775554087b8eee434b Mon Sep 17 00:00:00 2001
From: Mike Henry <11765982+mikemhenry@users.noreply.github.com>
Date: Fri, 14 Jul 2023 06:38:44 -0700
Subject: [PATCH] test for model fetch

---
 espaloma/utils/tests/test_model_fetch.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 espaloma/utils/tests/test_model_fetch.py

diff --git a/espaloma/utils/tests/test_model_fetch.py b/espaloma/utils/tests/test_model_fetch.py
new file mode 100644
index 00000000..6ab6dd58
--- /dev/null
+++ b/espaloma/utils/tests/test_model_fetch.py
@@ -0,0 +1,15 @@
+import espaloma as esp
+import torch
+from espaloma.utils.model_fetch import get_model
+from openff.toolkit.topology import Molecule
+
+def test_get_model(tmp_path):
+    model_dir = tmp_path / "latest"
+    model_path = get_model(model_dir=model_dir, disable_progress_bar=True)
+
+    molecule = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
+    molecule_graph = esp.Graph(molecule)
+
+    espaloma_model = torch.load(model_path)
+    espaloma_model.eval()
+    espaloma_model(molecule_graph.heterograph)