Require a test to be added for fast push

pkg/config/config.go

@@ -68,12 +68,17 @@ type Example struct {
 	Output string            `json:"output" yaml:"output"`
 }
 
+type Test struct {
+	Inputs map[string]string `json:"inputs" yaml:"inputs"`
+}
+
 type Config struct {
 	Build       *Build       `json:"build" yaml:"build"`
 	Image       string       `json:"image,omitempty" yaml:"image"`
 	Predict     string       `json:"predict,omitempty" yaml:"predict"`
 	Train       string       `json:"train,omitempty" yaml:"train"`
 	Concurrency *Concurrency `json:"concurrency,omitempty" yaml:"concurrency"`
+	Tests       []Test       `json:"tests,omitempty" yaml:"tests"`
 }
 
 func DefaultConfig() *Config {

pkg/config/config_test.go

@@ -693,3 +693,17 @@ torch==2.4.0
 torchvision==2.4.0`
 	require.Equal(t, expected, requirements)
 }
+
+func TestParseTests(t *testing.T) {
+	yamlString := `
+build:
+  run:
+  - command: "echo 'Hello, World!'"
+tests:
+- inputs:
+    s: world
+`
+	cfg, err := FromYAML([]byte(yamlString))
+	require.NoError(t, err)
+	require.Equal(t, map[string]string{"s": "world"}, cfg.Tests[0].Inputs)
+}

pkg/config/data/config_schema_v1.0.json

@@ -26,12 +26,18 @@
         },
         "python_version": {
           "$id": "#/properties/build/properties/python_version",
-          "type": ["string", "number"],
+          "type": [
+            "string",
+            "number"
+          ],
           "description": "The minor (`3.8`) or patch (`3.8.1`) version of Python to use."
         },
         "python_packages": {
           "$id": "#/properties/build/properties/python_packages",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of Python packages to install, in the format `package==version`.",
           "additionalItems": true,
           "items": {
@@ -46,7 +52,10 @@
         },
         "pre_install": {
           "$id": "#/properties/build/properties/pre_install",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of setup commands to run in the environment before your Python packages are installed.",
           "additionalItems": true,
           "items": {
@@ -66,7 +75,10 @@
         },
         "system_packages": {
           "$id": "#/properties/build/properties/system_packages",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of Ubuntu APT packages to install.",
           "additionalItems": true,
           "items": {
@@ -81,7 +93,10 @@
         },
         "run": {
           "$id": "#/properties/build/properties/run",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of setup commands to run in the environment after your system packages and Python packages have been installed. If you're familiar with Docker, it's like a `RUN` instruction in your `Dockerfile`.",
           "additionalItems": true,
           "items": {
@@ -105,7 +120,9 @@
                       "properties": {
                         "type": {
                           "type": "string",
-                          "enum": ["secret"]
+                          "enum": [
+                            "secret"
+                          ]
                         },
                         "id": {
                           "type": "string"
@@ -114,11 +131,17 @@
                           "type": "string"
                         }
                       },
-                      "required": ["type", "id", "target"]
+                      "required": [
+                        "type",
+                        "id",
+                        "target"
+                      ]
                     }
                   }
                 },
-                "required": ["command"]
+                "required": [
+                  "command"
+                ]
               }
             ]
           }
@@ -161,6 +184,19 @@
           "description": "The default target for number of concurrent predictions. This setting can be used by an autoscaler to determine when to scale a deployment of a model up or down."
         }
       }
+    },
+    "tests": {
+      "$id": "#/properties/tests",
+      "type": [
+        "array",
+        "null"
+      ],
+      "description": "A list of test cog commands to run.",
+      "additionalItems": false,
+      "items": {
+        "type": "object",
+        "additionalItems": true
+      }
     }
   },
   "additionalProperties": false

pkg/dockerfile/fast_generator.go

@@ -296,5 +296,8 @@ func (g *FastGenerator) validateConfig() error {
 	if len(g.Config.Build.Run) > 0 {
 		return errors.New("cog builds with --x-fast do not support build run commands.")
 	}
+	if len(g.Config.Tests) < 1 {
+		return errors.New("cog builds with --x-fast require at least 1 test.")
+	}
 	return nil
 }

pkg/dockerfile/fast_generator_test.go

@@ -18,6 +18,13 @@ func TestGenerate(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Inputs: map[string]string{
+					"s": "world",
+				},
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -56,6 +63,13 @@ func TestGenerateUVCacheMount(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Inputs: map[string]string{
+					"s": "world",
+				},
+			},
+		},
 	}
 	// Create matrix
 	matrix := MonobaseMatrix{
@@ -91,6 +105,13 @@ func TestGenerateCUDA(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Inputs: map[string]string{
+					"s": "world",
+				},
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -128,6 +149,13 @@ func TestGeneratePythonPackages(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Inputs: map[string]string{
+					"s": "world",
+				},
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -163,6 +191,13 @@ func TestGenerateVerboseEnv(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Inputs: map[string]string{
+					"s": "world",
+				},
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -198,6 +233,13 @@ func TestAptInstall(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Inputs: map[string]string{
+					"s": "world",
+				},
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -260,3 +302,34 @@ func TestValidateConfigWithBuildRunItems(t *testing.T) {
 	err = generator.validateConfig()
 	require.Error(t, err)
 }
+
+func TestValidateConfigWithNoTests(t *testing.T) {
+	dir := t.TempDir()
+	build := config.Build{
+		PythonVersion:  "3.8",
+		SystemPackages: []string{"git"},
+	}
+	config := config.Config{
+		Build: &build,
+	}
+	command := dockertest.NewMockCommand()
+	matrix := MonobaseMatrix{
+		Id:             1,
+		CudaVersions:   []string{"2.4"},
+		CudnnVersions:  []string{"1.0"},
+		PythonVersions: []string{"3.8"},
+		TorchVersions:  []string{"2.5.1"},
+		Venvs: []MonobaseVenv{
+			{
+				Python: "3.8",
+				Torch:  "2.5.1",
+				Cuda:   "2.4",
+			},
+		},
+	}
+	generator, err := NewFastGenerator(&config, dir, command, &matrix)
+	require.NoError(t, err)
+
+	err = generator.validateConfig()
+	require.Error(t, err)
+}

test-integration/test_integration/fixtures/fast-build/cog.yaml

@@ -7,3 +7,6 @@ build:
     - "git"
 predict: "predict.py:Predictor"
 image: "r8.im/<R8_USERNAME>/test"
+tests:
+- inputs:
+    s: world