encode_test.go

package jsonstat

import (
	"bytes"
	"encoding/json"
	"errors"
	"testing"
)

// This file tests the JSON-stat 2.0 wire-format encoder. Round-trip fidelity
// is the central guarantee: decode → encode → decode must produce a Dataset
// that is semantically equivalent to the original. We verify equivalence by
// re-walking the cube and comparing cell-by-cell, not by byte-equality (real
// fixtures vary in whitespace and key order).
//
// Beyond round-trip, each value form (dense/sparse), each status form
// (uniform/array/object), and the standalone-dimension and collection classes
// are exercised individually.

// ---------- Round-trip across canonical fixtures ----------

func TestEncodeRoundTripDatasets(t *testing.T) {
	// Every fixture in here must be a dataset-class document. (Collections
	// and standalone dimensions have their own tests below.)
	fixtures := []string{
		"oecd.json",
		"canada.json",
		"galicia.json",
		"hierarchy.json",
		"us-gsp.json",
		"us-unr.json",
		"order.json",
		"sparse.json",
		"status-array.json",
	}

	for _, name := range fixtures {
		t.Run(name, func(t *testing.T) {
			orig := mustDecode(t, loadFixture(t, name))
			origDS, err := orig.SingleDataset()
			if err != nil {
				t.Fatalf("fixture %s is not a single-dataset document: %v", name, err)
			}

			encoded, err := MarshalDataset(origDS)
			if err != nil {
				t.Fatalf("MarshalDataset(%s): %v", name, err)
			}

			// Sanity: encoded output must be valid JSON.
			if !json.Valid(encoded) {
				t.Fatalf("MarshalDataset(%s) produced invalid JSON: %s", name, encoded)
			}

			// Re-decode the encoded output.
			doc2, err := DecodeBytes(encoded)
			if err != nil {
				t.Fatalf("re-decode of encoded %s: %v\nencoded=%s", name, err, encoded)
			}
			ds2, err := doc2.SingleDataset()
			if err != nil {
				t.Fatalf("re-decoded doc has no dataset: %v", err)
			}

			assertDatasetsEquivalent(t, name, origDS, ds2)
		})
	}
}

// assertDatasetsEquivalent compares two Datasets by walking every cell. This
// is the strongest semantic check we can make without a generated equality
// method: the value/status stores, dimensions, and roles must all agree.
func assertDatasetsEquivalent(t *testing.T, label string, a, b *Dataset) {
	t.Helper()

	// Structural fields.
	if len(a.ID) != len(b.ID) {
		t.Fatalf("%s: ID length mismatch: %v vs %v", label, a.ID, b.ID)
	}
	for i := range a.ID {
		if a.ID[i] != b.ID[i] {
			t.Errorf("%s: ID[%d]: %q vs %q", label, i, a.ID[i], b.ID[i])
		}
	}
	if len(a.Size) != len(b.Size) {
		t.Fatalf("%s: Size length mismatch: %v vs %v", label, a.Size, b.Size)
	}
	for i := range a.Size {
		if a.Size[i] != b.Size[i] {
			t.Errorf("%s: Size[%d]: %d vs %d", label, i, a.Size[i], b.Size[i])
		}
	}
	if a.N() != b.N() {
		t.Fatalf("%s: N mismatch: %d vs %d", label, a.N(), b.N())
	}

	// Value form must round-trip (dense stays dense, sparse stays sparse).
	if a.ValueFormat() != b.ValueFormat() {
		t.Errorf("%s: ValueFormat changed: %s vs %s", label, a.ValueFormat(), b.ValueFormat())
	}
	// Status form must round-trip.
	if a.StatusFormat() != b.StatusFormat() {
		t.Errorf("%s: StatusFormat changed: %s vs %s", label, a.StatusFormat(), b.StatusFormat())
	}

	// Walk every cell and compare Value, Status, Missing.
	if a.N() == 0 {
		return
	}
	cellsA, _ := a.CellsSlice()
	cellsB, _ := b.CellsSlice()
	if len(cellsA) != len(cellsB) {
		t.Fatalf("%s: cell counts differ: %d vs %d", label, len(cellsA), len(cellsB))
	}
	mismatches := 0
	for i := range cellsA {
		if !cellsEqual(cellsA[i], cellsB[i]) {
			if mismatches < 5 {
				t.Errorf("%s: cell %d differs:\n  a=%#v\n  b=%#v", label, i, cellsA[i], cellsB[i])
			}
			mismatches++
		}
	}
	if mismatches > 0 {
		t.Errorf("%s: %d/%d cells differ", label, mismatches, len(cellsA))
	}
}

// ---------- Dense value round-trip ----------

func TestEncodeDenseValues(t *testing.T) {
	// oecd.json: dense 432-cell value array, status is a sparse object.
	ds := mustDataset(t, "oecd.json")
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	// Spot-check a value in the encoded JSON to confirm it's a dense array.
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(encoded, &probe); err != nil {
		t.Fatal(err)
	}
	valRaw := probe["value"]
	if len(valRaw) == 0 || valRaw[0] != '[' {
		t.Errorf("dense value should encode as array, got %s", valRaw)
	}

	// Check a specific known value.
	want := `5.943826289`
	if !bytes.Contains(valRaw, []byte(want)) {
		t.Errorf("encoded value missing %s: %s", want, valRaw)
	}
}

// ---------- Sparse value round-trip ----------

func TestEncodeSparseValues(t *testing.T) {
	// sparse.json: 4 keys {0,5,13,20}.
	ds := mustDataset(t, "sparse.json")
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(encoded, &probe); err != nil {
		t.Fatal(err)
	}
	valRaw := probe["value"]
	if len(valRaw) == 0 || valRaw[0] != '{' {
		t.Errorf("sparse value should encode as object, got %s", valRaw)
	}
	// Re-decode and confirm the four values survive.
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatal(err)
	}
	ds2 := doc2.Dataset
	c, _ := ds2.CellByFlat(13)
	if c.Value != 110 {
		t.Errorf("flat 13 after round-trip: got %v, want 110", c.Value)
	}
}

// ---------- Status forms ----------

func TestEncodeStatusUniform(t *testing.T) {
	// canada.json has a string status. After round-trip, status form must
	// still be "uniform".
	ds := mustDataset(t, "canada.json")
	if ds.StatusFormat() != "uniform" {
		t.Fatalf("canada status form pre-encode: got %s, want uniform", ds.StatusFormat())
	}
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatal(err)
	}
	if got := doc2.Dataset.StatusFormat(); got != "uniform" {
		t.Errorf("canada status form post-encode: got %s, want uniform", got)
	}
}

func TestEncodeStatusArray(t *testing.T) {
	ds := mustDataset(t, "status-array.json")
	if ds.StatusFormat() != "array" {
		t.Fatalf("status-array form pre-encode: got %s, want array", ds.StatusFormat())
	}
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatal(err)
	}
	if got := doc2.Dataset.StatusFormat(); got != "array" {
		t.Errorf("status-array form post-encode: got %s, want array", got)
	}
	// Spot-check a status value at flat 4 ("p").
	c, _ := doc2.Dataset.CellByFlat(4)
	if c.Status != "p" {
		t.Errorf("flat 4 status after round-trip: got %q, want 'p'", c.Status)
	}
}

func TestEncodeStatusObject(t *testing.T) {
	// oecd.json has object-form status.
	ds := mustDataset(t, "oecd.json")
	if ds.StatusFormat() != "object" {
		t.Fatalf("oecd status form pre-encode: got %s, want object", ds.StatusFormat())
	}
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatal(err)
	}
	if got := doc2.Dataset.StatusFormat(); got != "object" {
		t.Errorf("oecd status form post-encode: got %s, want object", got)
	}
}

func TestEncodeStatusAbsent(t *testing.T) {
	// us-gsp.json has no status property.
	ds := mustDataset(t, "us-gsp.json")
	if ds.StatusFormat() != "none" {
		t.Fatalf("us-gsp status form: got %s, want none", ds.StatusFormat())
	}
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(encoded, &probe); err != nil {
		t.Fatal(err)
	}
	if _, present := probe["status"]; present {
		t.Errorf("us-gsp encoded a status property, want none")
	}
}

// ---------- Hierarchy ----------

func TestEncodeHierarchy(t *testing.T) {
	ds := mustDataset(t, "hierarchy.json")
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatal(err)
	}
	ds2 := doc2.Dataset
	// Find the dimension with a hierarchy and confirm it survived.
	anyHierarchy := false
	for _, id := range ds.ID {
		if ds.Dimensions[id].HasHierarchy() {
			anyHierarchy = true
			if !ds2.Dimensions[id].HasHierarchy() {
				t.Errorf("dim %s: hierarchy lost on round-trip", id)
			}
		}
	}
	if !anyHierarchy {
		t.Errorf("no hierarchy found in source dataset")
	}
}

// ---------- Roles ----------

func TestEncodeRoles(t *testing.T) {
	// oecd.json has role {time:[year], geo:[area], metric:[concept]}.
	ds := mustDataset(t, "oecd.json")
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatal(err)
	}
	ds2 := doc2.Dataset
	for _, r := range []Role{RoleTime, RoleGeo, RoleMetric} {
		if !ds2.HasRole(r) {
			t.Errorf("role %s lost on round-trip", r)
		}
	}
	// Verify role-to-dim mapping.
	for _, r := range []Role{RoleTime, RoleGeo, RoleMetric} {
		got := ds2.Role[r]
		want := ds.Role[r]
		if len(got) != len(want) {
			t.Errorf("role %s mapping len: got %d, want %d", r, len(got), len(want))
			continue
		}
		for i := range got {
			if got[i] != want[i] {
				t.Errorf("role %s[%d]: got %q, want %q", r, i, got[i], want[i])
			}
		}
	}
}

// ---------- Collection ----------

func TestEncodeCollection(t *testing.T) {
	doc := mustDecode(t, loadFixture(t, "collection.json"))
	if !doc.IsCollection() {
		t.Fatalf("collection.json is not a collection")
	}
	encoded, err := MarshalCollection(doc.Collection)
	if err != nil {
		t.Fatal(err)
	}
	if !json.Valid(encoded) {
		t.Fatalf("encoded collection is invalid JSON: %s", encoded)
	}
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatalf("re-decode: %v", err)
	}
	if !doc2.IsCollection() {
		t.Fatalf("re-decoded doc is not a collection")
	}
	if got := len(doc2.Collection.Items); got != len(doc.Collection.Items) {
		t.Errorf("items count: got %d, want %d", got, len(doc.Collection.Items))
	}
}

// ---------- Standalone dimension ----------

func TestEncodeStandaloneDimension(t *testing.T) {
	doc := mustDecode(t, loadFixture(t, "dimension-only.json"))
	if !doc.IsDimension() {
		t.Fatalf("dimension-only.json is not a standalone dimension")
	}
	encoded, err := MarshalDimension(doc.Dimension)
	if err != nil {
		t.Fatal(err)
	}
	if !json.Valid(encoded) {
		t.Fatalf("encoded dimension is invalid JSON: %s", encoded)
	}
	doc2, err := DecodeBytes(encoded)
	if err != nil {
		t.Fatalf("re-decode: %v", err)
	}
	if !doc2.IsDimension() {
		t.Fatalf("re-decoded doc is not a standalone dimension")
	}
	// Verify the role survived.
	if doc2.Dimension.Role != doc.Dimension.Role {
		t.Errorf("dim role: got %q, want %q", doc2.Dimension.Role, doc.Dimension.Role)
	}
}

// ---------- Pre-2.0 bundle rejection ----------

func TestEncodeBundleRejected(t *testing.T) {
	doc := mustDecode(t, loadFixture(t, "bundle-pre2.0.json"))
	if !doc.IsBundle {
		t.Fatalf("bundle-pre2.0.json is not a bundle")
	}
	_, err := MarshalDocument(doc)
	if err == nil {
		t.Fatalf("MarshalDocument on bundle: want error, got nil")
	}
	if !errors.Is(err, ErrUnsupportedVersion) {
		t.Errorf("MarshalDocument on bundle: got %v, want ErrUnsupportedVersion", err)
	}
}

// ---------- MarshalJSON interfaces ----------

func TestDatasetMarshalJSON(t *testing.T) {
	ds := mustDataset(t, "oecd.json")
	b, err := ds.MarshalJSON()
	if err != nil {
		t.Fatal(err)
	}
	if !json.Valid(b) {
		t.Fatalf("Dataset.MarshalJSON produced invalid JSON")
	}
}

func TestEncodeBytes(t *testing.T) {
	doc := mustDecode(t, loadFixture(t, "oecd.json"))
	b, err := EncodeBytes(doc)
	if err != nil {
		t.Fatal(err)
	}
	if !json.Valid(b) {
		t.Fatalf("EncodeBytes produced invalid JSON")
	}
}

func TestEncodeToBuffer(t *testing.T) {
	doc := mustDecode(t, loadFixture(t, "status-array.json"))
	var buf bytes.Buffer
	if err := Encode(&buf, doc); err != nil {
		t.Fatal(err)
	}
	if !json.Valid(buf.Bytes()) {
		t.Fatalf("Encode produced invalid JSON")
	}
}

// ---------- Nil handling ----------

func TestEncodeNilDataset(t *testing.T) {
	_, err := MarshalDataset(nil)
	if err == nil {
		t.Fatalf("MarshalDataset(nil): want error")
	}
}

func TestEncodeNilDocument(t *testing.T) {
	_, err := MarshalDocument(nil)
	if err == nil {
		t.Fatalf("MarshalDocument(nil): want error")
	}
}

func TestEncodeEmptyDocument(t *testing.T) {
	doc := &Document{}
	_, err := MarshalDocument(doc)
	if err == nil {
		t.Fatalf("MarshalDocument(empty): want error")
	}
}

// ---------- Property order check ----------

func TestEncodeKeyOrder(t *testing.T) {
	// The canonical JSON-stat property order for a dataset starts with
	// "class". This is a smoke check that the ordered writer works.
	ds := mustDataset(t, "oecd.json")
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	// The first key after the opening brace must be "class".
	prefix := `{"class":`
	if !bytes.HasPrefix(encoded, []byte(prefix)) {
		t.Errorf("encoded dataset does not start with %q: got %s", prefix, encoded[:40])
	}
}

// ---------- Category index form: array ----------

func TestEncodeCategoryIndexArray(t *testing.T) {
	// The encoder always emits category.index as an array. Re-decoding must
	// work because decodeCategoryIndex accepts arrays.
	ds := mustDataset(t, "oecd.json")
	encoded, err := MarshalDataset(ds)
	if err != nil {
		t.Fatal(err)
	}
	var probe struct {
		Dimension map[string]struct {
			Category struct {
				Index []string `json:"index"`
			} `json:"category"`
		} `json:"dimension"`
	}
	if err := json.Unmarshal(encoded, &probe); err != nil {
		t.Fatal(err)
	}
	areaDim := probe.Dimension["area"]
	if len(areaDim.Category.Index) == 0 {
		t.Errorf("area dimension category.index is not an array")
	}
	// The area dimension in oecd.json has size 36.
	if len(areaDim.Category.Index) != 36 {
		t.Errorf("area category.index length: got %d, want 36", len(areaDim.Category.Index))
	}
}