context_test.go 614 B

1234567891011121314151617181920212223242526272829
  1. //go:build integration
  2. package integration
  3. import (
  4. "context"
  5. "net/http"
  6. "testing"
  7. "time"
  8. "ollama.com/api"
  9. )
  10. func TestContextExhaustion(t *testing.T) {
  11. ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) // TODO maybe shorter?
  12. defer cancel()
  13. // Set up the test data
  14. req := api.GenerateRequest{
  15. Model: "llama2",
  16. Prompt: "Write me a story with a ton of emojis?",
  17. Stream: &stream,
  18. Options: map[string]interface{}{
  19. "temperature": 0,
  20. "seed": 123,
  21. "num_ctx": 128,
  22. },
  23. }
  24. GenerateTestHelper(ctx, t, &http.Client{}, req, []string{"once", "upon", "lived"})
  25. }