add support for automatically unloading a model (#10) (#14)

* Make starting upstream process on-demand (#10)
* Add automatic unload of model after TTL is reached
* add `ttl` configuration parameter to models in seconds, default is 0 (never unload)

proxy/process_test.go

@@ -2,14 +2,17 @@ package proxy
 
 import (
 	"fmt"
+	"io"
 	"math/rand"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"runtime"
-	"strings"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 // Check if the binary exists
@@ -29,53 +32,40 @@ func getBinaryPath() string {
 	return filepath.Join("..", "build", fmt.Sprintf("simple-responder_%s_%s", goos, goarch))
 }
 
-func TestProcess_ProcessStartStop(t *testing.T) {
+func getTestSimpleResponderConfig(expectedMessage string) ModelConfig {
 	// Define the range
 	min := 12000
 	max := 13000
 
 	// Generate a random number between 12000 and 13000
 	randomPort := rand.Intn(max-min+1) + min
-
 	binaryPath := getBinaryPath()
 
-	// Create a log monitor
-	logMonitor := NewLogMonitor()
-
-	expectedMessage := "testing91931"
-
 	// Create a process configuration
-	config := ModelConfig{
+	return ModelConfig{
 		Cmd:           fmt.Sprintf("%s --port %d --respond '%s'", binaryPath, randomPort, expectedMessage),
 		Proxy:         fmt.Sprintf("http://127.0.0.1:%d", randomPort),
 		CheckEndpoint: "/health",
 	}
+}
+
+func TestProcess_AutomaticallyStartsUpstream(t *testing.T) {
+	logMonitor := NewLogMonitorWriter(io.Discard)
+	expectedMessage := "testing91931"
+	config := getTestSimpleResponderConfig(expectedMessage)
 
 	// Create a process
-	process := NewProcess("test-process", config, logMonitor)
-
-	// Start the process
-	t.Logf("Starting %s on port %d", binaryPath, randomPort)
-	err := process.Start(5)
-	if err != nil {
-		t.Fatalf("Failed to start process: %v", err)
-	}
-
-	// Create a test request
+	process := NewProcess("test-process", 5, config, logMonitor)
 	req := httptest.NewRequest("GET", "/", nil)
 	w := httptest.NewRecorder()
 
-	// Proxy the request
+	// process is automatically started
+	assert.False(t, process.IsRunning())
 	process.ProxyRequest(w, req)
+	assert.True(t, process.IsRunning())
 
-	// Check the response
-	if w.Code != http.StatusOK {
-		t.Errorf("Expected status code %d, got %d", http.StatusOK, w.Code)
-	}
-
-	if !strings.Contains(w.Body.String(), expectedMessage) {
-		t.Errorf("Expected body to contain '%s', got %q", expectedMessage, w.Body.String())
-	}
+	assert.Equal(t, http.StatusOK, w.Code, "Expected status code %d, got %d", http.StatusOK, w.Code)
+	assert.Contains(t, w.Body.String(), expectedMessage)
 
 	// Stop the process
 	process.Stop()
@@ -86,8 +76,53 @@ func TestProcess_ProcessStartStop(t *testing.T) {
 	// Proxy the request
 	process.ProxyRequest(w, req)
 
-	// Check the response
-	if w.Code == http.StatusInternalServerError {
-		t.Errorf("Expected status code %d, got %d", http.StatusInternalServerError, w.Code)
+	// should have automatically started the process again
+	if w.Code != http.StatusOK {
+		t.Errorf("Expected status code %d, got %d", http.StatusOK, w.Code)
 	}
 }
+
+// test that the automatic start returns the expected error type
+func TestProcess_BrokenModelConfig(t *testing.T) {
+	// Create a process configuration
+	config := ModelConfig{
+		Cmd:           "nonexistant-command",
+		Proxy:         "http://127.0.0.1:9913",
+		CheckEndpoint: "/health",
+	}
+
+	process := NewProcess("broken", 1, config, NewLogMonitor())
+	req := httptest.NewRequest("GET", "/", nil)
+	w := httptest.NewRecorder()
+	process.ProxyRequest(w, req)
+	assert.Equal(t, http.StatusInternalServerError, w.Code)
+	assert.Contains(t, w.Body.String(), "unable to start process")
+}
+
+func TestProcess_UnloadAfterTTL(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping long auto unload TTL test")
+	}
+
+	expectedMessage := "I_sense_imminent_danger"
+	config := getTestSimpleResponderConfig(expectedMessage)
+	assert.Equal(t, 0, config.UnloadAfter)
+	config.UnloadAfter = 3 // seconds
+	assert.Equal(t, 3, config.UnloadAfter)
+
+	process := NewProcess("ttl", 2, config, NewLogMonitorWriter(io.Discard))
+	req := httptest.NewRequest("GET", "/", nil)
+	w := httptest.NewRecorder()
+
+	// Proxy the request (auto start)
+	process.ProxyRequest(w, req)
+	assert.Equal(t, http.StatusOK, w.Code, "Expected status code %d, got %d", http.StatusOK, w.Code)
+	assert.Contains(t, w.Body.String(), expectedMessage)
+
+	assert.True(t, process.IsRunning())
+
+	// wait 5 seconds
+	time.Sleep(5 * time.Second)
+
+	assert.False(t, process.IsRunning())
+}