Merge pull request #5 from viam-labs/model-ext-check

check model extension for onnx

.github/workflows/test.yml

@@ -19,11 +19,12 @@ jobs:
     - name: Install dependencies
       run: go mod download
 
-    - name: Run unit tests
-      run: make test
-
     - name: golangci-lint
       uses: golangci/golangci-lint-action@v6
       with:
         version: v1.59
-        args: --disable errcheck --timeout 10m
+        args: --timeout 10m
+
+    - name: Run unit tests
+      run: make test
+

Makefile

@@ -3,7 +3,7 @@ MOD_OS := $(shell uname -s)
 test:
 	go test
 lint:
-	golangci-lint run --disable errcheck
+	golangci-lint run
 module.tar.gz:
 ifeq ($(MOD_OS),Darwin)
 ifeq ($(MOD_ARCH),x86_64)
@@ -60,4 +60,4 @@ third_party/onnx-android-$(SO_ARCH).so: onnxruntime-android-$(ONNX_VERSION).aar
 	cp jni/$(SO_ARCH)/libonnxruntime.so $@
 
 bundle-droid-$(SO_ARCH).tar.gz: module third_party/onnx-android-$(SO_ARCH).so
-	tar -czf $@ $^
+	tar -czf $@ $^

cmd/module/main.go

@@ -36,7 +36,7 @@ func realMain() error {
 	if err != nil {
 		return err
 	}
-	defer onnxClose(ctx)
+	defer onnxClose(ctx) //nolint:errcheck
 	<-ctx.Done()
 	return nil
 }

onnx_cpu.go

@@ -3,6 +3,7 @@ package onnx_cpu
 
 import (
 	"context"
+	"path"
 	"runtime"
 	"strings"
 
@@ -48,9 +49,14 @@ type Config struct {
 }
 
 // Validate makes sure that the required model path is not empty
-func (cfg *Config) Validate(path string) ([]string, error) {
+func (cfg *Config) Validate(validatePath string) ([]string, error) {
 	if cfg.ModelPath == "" {
-		return nil, utils.NewConfigValidationFieldRequiredError(path, "model_path")
+		return nil, utils.NewConfigValidationFieldRequiredError(validatePath, "model_path")
+	}
+	ext := path.Ext(cfg.ModelPath)
+	if ext != ".onnx" {
+		base := path.Base(cfg.ModelPath)
+		return nil, errors.Errorf("model_path filename must end in .onnx. The filename is %s", base)
 	}
 	return nil, nil
 }

onnx_cpu_test.go

@@ -110,3 +110,22 @@ func TestImageDetection(t *testing.T) {
 	err = theModel.Close(context.Background())
 	test.That(t, err, test.ShouldBeNil)
 }
+
+func TestValidate(t *testing.T) {
+	// correct
+	cfg := &Config{"./test_files/ir_mobilenet.onnx", "/path/to/labels.txt"}
+	deps, err := cfg.Validate("")
+	test.That(t, err, test.ShouldBeNil)
+	test.That(t, deps, test.ShouldBeNil)
+	// empty
+	cfg = &Config{"", "/path/to/labels.txt"}
+	_, err = cfg.Validate("")
+	test.That(t, err, test.ShouldNotBeNil)
+	test.That(t, err.Error(), test.ShouldContainSubstring, "model_path")
+	// incorrect
+	cfg = &Config{"/path/to/other_model.tflite", "/path/to/labels.txt"}
+	_, err = cfg.Validate("")
+	test.That(t, err, test.ShouldNotBeNil)
+	test.That(t, err.Error(), test.ShouldContainSubstring, "must end in .onnx")
+
+}