From 8de18bc1d51226a13618fb25cd20e146b48408ff Mon Sep 17 00:00:00 2001
From: Joshua Lochner <admin@xenova.com>
Date: Tue, 2 May 2023 05:37:39 +0200
Subject: [PATCH] Create hub.test.js

---
 tests/hub.test.js | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)
 create mode 100644 tests/hub.test.js

diff --git a/tests/hub.test.js b/tests/hub.test.js
new file mode 100644
index 0000000..50bd29c
--- /dev/null
+++ b/tests/hub.test.js
@@ -0,0 +1,37 @@
+
+
+import { env, AutoModel, PreTrainedModel } from '../src/transformers.js';
+
+// env.allowRemoteModels = false;
+const MAX_TEST_EXECUTION_TIME = 60_000; // 60 seconds
+
+// TODO: Set cache folder to a temp directory
+
+describe('Hub', () => {
+
+    describe('Loading models', () => {
+
+        it('should load a model from the local cache', async () => {
+            // 1. Local model exists (doesn't matter about status of remote file since local is tried first)
+            let model = await AutoModel.from_pretrained('t5-small');
+            expect(model).toBeInstanceOf(PreTrainedModel);
+        }, MAX_TEST_EXECUTION_TIME);
+
+        it('should load a model from the remote cache', async () => {
+            // 2. Local model doesn't exist, remote file exists
+            // This tests that fallback functionality is working
+            let model = await AutoModel.from_pretrained('Xenova/t5-small-web');
+            expect(model).toBeInstanceOf(PreTrainedModel);
+        }, MAX_TEST_EXECUTION_TIME);
+
+        it('should fail to load a model', async () => {
+            // 3. Local model doesn't exist, remote file doesn't exist
+            // This tests that error handling is working.
+            await expect(
+                AutoModel.from_pretrained('Xenova/this-model-does-not-exist')
+            ).rejects
+                .toBeInstanceOf(Error);
+        }, MAX_TEST_EXECUTION_TIME);
+    });
+
+});