cmd/python-provider/main_test.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"

	spine "pysolate-spine"
)

func TestMain(m *testing.M) {
	if err := os.Chdir(filepath.Join("..", "..")); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	os.Exit(m.Run())
}

func TestPythonProviderGuest(t *testing.T) {
	runner := newTestRunner(t)
	out, err := runner.Run(context.Background(), readGuest(t), map[string]any{"numbers": []int{2, 4, 6, 8}})
	if err != nil {
		t.Fatal(err)
	}
	var got map[string]any
	if err := json.Unmarshal(out.Value, &got); err != nil {
		t.Fatal(err)
	}
	want := map[string]any{"count": float64(4), "total": float64(20), "mean": float64(5)}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("got %v, want %v", got, want)
	}
}

func TestPythonProviderRejectsBadInput(t *testing.T) {
	runner := newTestRunner(t)
	_, err := runner.Run(context.Background(), readGuest(t), map[string]any{"numbers": "not a list"})
	if err == nil || !strings.Contains(err.Error(), "report_stats input:") {
		t.Fatalf("expected an input validation error, got %v", err)
	}
}

func TestPythonProviderWorkerFailure(t *testing.T) {
	runner := newTestRunner(t)
	dir := t.TempDir()
	fakePython := filepath.Join(dir, "python3")
	if err := os.WriteFile(fakePython, []byte("#!/bin/sh\necho worker-broke >&2\nexit 7\n"), 0700); err != nil {
		t.Fatal(err)
	}
	oldPath := os.Getenv("PATH")
	if err := os.Setenv("PATH", dir); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = os.Setenv("PATH", oldPath) })

	_, err := runner.Run(context.Background(), readGuest(t), map[string]any{"numbers": []int{1}})
	if err == nil || !strings.Contains(err.Error(), "report_stats worker failed") || !strings.Contains(err.Error(), "worker-broke") {
		t.Fatalf("expected surfaced worker failure, got %v", err)
	}
}

func TestPythonProviderMalformedWorkerOutput(t *testing.T) {
	runner := newTestRunner(t)
	dir := t.TempDir()
	fakePython := filepath.Join(dir, "python3")
	if err := os.WriteFile(fakePython, []byte("#!/bin/sh\nprintf 'not-json\\n'\n"), 0700); err != nil {
		t.Fatal(err)
	}
	oldPath := os.Getenv("PATH")
	if err := os.Setenv("PATH", dir); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = os.Setenv("PATH", oldPath) })

	_, err := runner.Run(context.Background(), readGuest(t), map[string]any{"numbers": []int{1}})
	if err == nil || !strings.Contains(err.Error(), "worker returned malformed JSON") {
		t.Fatalf("expected malformed worker JSON error, got %v", err)
	}
}

func TestReportWorkerSubprocess(t *testing.T) {
	cases := []struct {
		name, input string
		wantCode    int
		wantOutput  string
	}{
		{"valid", `{"numbers":[2,4,6]}`, 0, `{"count": 3, "total": 12}`},
		{"malformed", `{"numbers":`, 2, ""},
		{"wrong shape", `{"numbers":[true]}`, 2, ""},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			command := exec.Command("python3", "providers/report.py")
			command.Stdin = strings.NewReader(tc.input)
			output, err := command.CombinedOutput()
			code := 0
			if err != nil {
				if exit, ok := err.(*exec.ExitError); ok {
					code = exit.ExitCode()
				} else {
					t.Fatal(err)
				}
			}
			if code != tc.wantCode {
				t.Fatalf("exit code %d, want %d (output: %s)", code, tc.wantCode, output)
			}
			if tc.wantOutput != "" && !bytes.Contains(output, []byte(tc.wantOutput)) {
				t.Fatalf("output %q does not contain %q", output, tc.wantOutput)
			}
		})
	}
}

func newTestRunner(t *testing.T) *spine.Runner {
	t.Helper()
	wasm, err := os.ReadFile("dist/spine.wasm")
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	t.Cleanup(cancel)
	runner, err := spine.New(ctx, wasm, spine.Manifest{
		"report_stats": {Call: reportStats},
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = runner.Close(context.Background()) })
	return runner
}

func readGuest(t *testing.T) string {
	t.Helper()
	code, err := os.ReadFile("examples/python_provider.py")
	if err != nil {
		t.Fatal(err)
	}
	return string(code)
}