code.oscarkilo.com/klex-git/api/messages_test.go

..
api.go
datasets.go
embed.go
f.go
fake.go
fake_test.go
funcs.go
funcs_test.go
klex.go
messages.go
messages_test.go
pipelines.go
stream.go
worker.go
package api

// These tests pin the wire format of ContentBlock so callers across
// //funky, //ithfm, and other consumers can rely on stable JSON
// shapes. They fail loudly if the JSON tags or field names ever
// drift.

import "encoding/json"
import "strings"
import "testing"

func TestDocumentBlockMarshal(t *testing.T) {
  block := ContentBlock{
    Type: "document",
    Source: &ContentSource{
      Type: "base64",
      MediaType: "application/pdf",
      Data: "JVBERi0xLjQK",
    },
  }
  buf, err := json.Marshal(block)
  if err != nil {
    t.Fatalf("marshal: %v", err)
  }
  want := `{"type":"document","source":{"type":"base64","media_type":"application/pdf","data":"JVBERi0xLjQK"}}`
  if string(buf) != want {
    t.Errorf(
      "marshal mismatch\nwant: %s\n got: %s", want, buf)
  }
}

func TestDocumentBlockUnmarshal(t *testing.T) {
  in := []byte(
    `{"type":"document","source":{"type":"base64",` +
      `"media_type":"application/pdf","data":"JVBERi0xLjQK"}}`)
  var block ContentBlock
  if err := json.Unmarshal(in, &block); err != nil {
    t.Fatalf("unmarshal: %v", err)
  }
  if block.Type != "document" {
    t.Errorf("Type = %q, want %q", block.Type, "document")
  }
  if block.Source == nil {
    t.Fatalf("Source is nil")
  }
  if block.Source.MediaType != "application/pdf" {
    t.Errorf(
      "MediaType = %q, want application/pdf",
      block.Source.MediaType)
  }
  if block.Source.Data != "JVBERi0xLjQK" {
    t.Errorf("Data = %q", block.Source.Data)
  }
}

func TestDocumentBlockImageMimeRoundTrip(t *testing.T) {
  // A "document" with an image/* MIME is how the wire format
  // carries what used to live under Type:"image".
  block := ContentBlock{
    Type: "document",
    Source: &ContentSource{
      Type: "base64",
      MediaType: "image/png",
      Data: "iVBORw0KGgo=",
    },
  }
  buf, err := json.Marshal(block)
  if err != nil {
    t.Fatalf("marshal: %v", err)
  }
  if !strings.Contains(string(buf), `"type":"document"`) {
    t.Errorf("expected type:document, got %s", buf)
  }
  if !strings.Contains(string(buf), `"media_type":"image/png"`) {
    t.Errorf("expected media_type image/png, got %s", buf)
  }
  var got ContentBlock
  if err := json.Unmarshal(buf, &got); err != nil {
    t.Fatalf("unmarshal: %v", err)
  }
  if got.Type != "document" {
    t.Errorf("Type = %q", got.Type)
  }
  if got.Source == nil ||
    got.Source.MediaType != "image/png" {
    t.Errorf("Source mismatch: %+v", got.Source)
  }
}

func TestTextBlockUnchanged(t *testing.T) {
  // Documents the unchanged text-block wire format.
  block := ContentBlock{Type: "text", Text: "hello"}
  buf, err := json.Marshal(block)
  if err != nil {
    t.Fatalf("marshal: %v", err)
  }
  if string(buf) != `{"type":"text","text":"hello"}` {
    t.Errorf("unexpected text-block JSON: %s", buf)
  }
}

func TestNewDocumentBlock(t *testing.T) {
  png := []byte{
    0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A,
    // padding so DetectContentType has enough to work with
    0, 0, 0, 0, 0, 0, 0, 0,
  }
  blk := NewDocumentBlock(png)
  if blk.Type != "document" {
    t.Errorf("Type = %q, want document", blk.Type)
  }
  if blk.Source == nil {
    t.Fatal("Source is nil")
  }
  if blk.Source.Type != "base64" {
    t.Errorf("Source.Type = %q, want base64",
      blk.Source.Type)
  }
  if blk.Source.MediaType != "image/png" {
    t.Errorf("MediaType = %q, want image/png",
      blk.Source.MediaType)
  }
  // The encoded bytes round-trip back to the original.
  buf, err := json.Marshal(blk)
  if err != nil {
    t.Fatalf("marshal: %v", err)
  }
  if !strings.Contains(string(buf), `"type":"document"`) {
    t.Errorf("marshaled JSON missing type:document: %s", buf)
  }
}