From aeb1fb5192be2e0990800f52183f3bd7cf15be54 Mon Sep 17 00:00:00 2001
From: Daniel Hiltgen <daniel@ollama.com>
Date: Thu, 4 Apr 2024 07:42:17 -0700
Subject: [PATCH] Add test case for context exhaustion

Confirmed this fails on 0.1.30 with known regression
but passes on main
---
 integration/context_test.go | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100644 integration/context_test.go

diff --git a/integration/context_test.go b/integration/context_test.go
new file mode 100644
index 00000000..80ea540b
--- /dev/null
+++ b/integration/context_test.go
@@ -0,0 +1,29 @@
+//go:build integration
+
+package integration
+
+import (
+	"context"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/ollama/ollama/api"
+)
+
+func TestContextExhaustion(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) // TODO maybe shorter?
+	defer cancel()
+	// Set up the test data
+	req := api.GenerateRequest{
+		Model:  "llama2",
+		Prompt: "Write me a story with a ton of emojis?",
+		Stream: &stream,
+		Options: map[string]interface{}{
+			"temperature": 0,
+			"seed":        123,
+			"num_ctx":     128,
+		},
+	}
+	GenerateTestHelper(ctx, t, &http.Client{}, req, []string{"once", "upon", "lived"})
+}