Support parameters for fetching models

Author: erbridge

lib/ReplicateClient.js

@@ -21,14 +21,33 @@ export default class ReplicateClient {
     }
   }
 
-  model(fullId) {
-    const idComponents = /^([\w-]+)\/([\w-]+)@(\w+)$/.exec(fullId);
+  model(idOrParams) {
+    let owner, name, version;
 
-    if (!idComponents) {
-      throw new ReplicateError(`Invalid ID: ${fullId}`);
-    }
+    if (typeof idOrParams === "string") {
+      const idComponents = /^([\w-]+)\/([\w-]+)@(\w+)$/.exec(idOrParams);
+
+      if (!idComponents) {
+        throw new ReplicateError(`Invalid ID: ${idOrParams}`);
+      }
 
-    const [, owner, name, version] = idComponents;
+      [, owner, name, version] = idComponents;
+    } else {
+      if (idOrParams.owner) {
+        owner = idOrParams.owner;
+        name = idOrParams.name;
+      } else {
+        const nameComponents = /^([\w-]+)\/([\w-]+)$/.exec(idOrParams.name);
+
+        if (!nameComponents) {
+          throw new ReplicateError(`Invalid name: ${idOrParams.name}`);
+        }
+
+        [, owner, name] = nameComponents;
+      }
+
+      version = idOrParams.version;
+    }
 
     return new Model({ owner, name, version }, this);
   }

lib/ReplicateClient.test.js

@@ -71,6 +71,37 @@ describe("model()", () => {
 
     expect(model).toBeInstanceOf(Model);
   });
+
+  it("parses ID string", () => {
+    const model = client.model("test-owner/test-name@testversion");
+
+    expect(model.owner).toBe("test-owner");
+    expect(model.name).toBe("test-name");
+    expect(model.version).toBe("testversion");
+  });
+
+  it("handles complete params", () => {
+    const model = client.model({
+      owner: "test-owner",
+      name: "test-name",
+      version: "testversion",
+    });
+
+    expect(model.owner).toBe("test-owner");
+    expect(model.name).toBe("test-name");
+    expect(model.version).toBe("testversion");
+  });
+
+  it("handles params with combined name", () => {
+    const model = client.model({
+      name: "test-owner/test-name",
+      version: "testversion",
+    });
+
+    expect(model.owner).toBe("test-owner");
+    expect(model.name).toBe("test-name");
+    expect(model.version).toBe("testversion");
+  });
 });
 
 describe("prediction()", () => {