From 239c3acd1d6d398b8d64280e3854f37aa439f89a Mon Sep 17 00:00:00 2001
From: Nic Cope <nicc@rk0n.org>
Date: Sun, 1 Oct 2023 17:11:23 -0700
Subject: [PATCH] Add tests for protobuf map fields

This switches to using a dedicated test.proto file, since the
descriptor.proto file previously being used doesn't have any map fields.

Signed-off-by: Nic Cope <nicc@rk0n.org>
---
 starlark/eval_test.go              | 28 +++++++++++++++++----
 starlark/testdata/proto.star       | 40 ++++++++++++++++++++++++------
 starlark/testdata/proto/test.fds   | 13 ++++++++++
 starlark/testdata/proto/test.proto | 11 ++++++++
 4 files changed, 79 insertions(+), 13 deletions(-)
 create mode 100644 starlark/testdata/proto/test.fds
 create mode 100644 starlark/testdata/proto/test.proto

diff --git a/starlark/eval_test.go b/starlark/eval_test.go
index 66786711..581f6bd3 100644
--- a/starlark/eval_test.go
+++ b/starlark/eval_test.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"os"
 	"os/exec"
 	"path/filepath"
 	"reflect"
@@ -19,15 +20,16 @@ import (
 	"go.starlark.net/internal/chunkedfile"
 	"go.starlark.net/lib/json"
 	starlarkmath "go.starlark.net/lib/math"
-	"go.starlark.net/lib/proto"
+	starlarkproto "go.starlark.net/lib/proto"
 	"go.starlark.net/lib/time"
 	"go.starlark.net/starlark"
 	"go.starlark.net/starlarkstruct"
 	"go.starlark.net/starlarktest"
 	"go.starlark.net/syntax"
-	"google.golang.org/protobuf/reflect/protoregistry"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protodesc"
 
-	_ "google.golang.org/protobuf/types/descriptorpb" // example descriptor needed for lib/proto tests
+	"google.golang.org/protobuf/types/descriptorpb"
 )
 
 // A test may enable non-standard options by containing (e.g.) "option:recursion".
@@ -115,7 +117,23 @@ func TestExecFile(t *testing.T) {
 	testdata := starlarktest.DataFile("starlark", ".")
 	thread := &starlark.Thread{Load: load}
 	starlarktest.SetReporter(thread, t)
-	proto.SetPool(thread, protoregistry.GlobalFiles)
+
+	// This proto is used for the proto.star tests. It's generated by running:
+	// protoc --descriptor_set_out=test.fds test.proto
+	data, err := os.ReadFile("testdata/proto/test.fds")
+	if err != nil {
+		t.Fatal(err)
+	}
+	fds := &descriptorpb.FileDescriptorSet{}
+	if err := proto.Unmarshal(data, fds); err != nil {
+		t.Fatal(err)
+	}
+	pool, err := protodesc.NewFiles(fds)
+	if err != nil {
+		t.Fatal(err)
+	}
+	starlarkproto.SetPool(thread, pool)
+
 	for _, file := range []string{
 		"testdata/assign.star",
 		"testdata/bool.star",
@@ -207,7 +225,7 @@ func load(thread *starlark.Thread, module string) (starlark.StringDict, error) {
 		return starlark.StringDict{"math": starlarkmath.Module}, nil
 	}
 	if module == "proto.star" {
-		return starlark.StringDict{"proto": proto.Module}, nil
+		return starlark.StringDict{"proto": starlarkproto.Module}, nil
 	}
 
 	// TODO(adonovan): test load() using this execution path.
diff --git a/starlark/testdata/proto.star b/starlark/testdata/proto.star
index ed91716c..fb4a38bc 100644
--- a/starlark/testdata/proto.star
+++ b/starlark/testdata/proto.star
@@ -3,14 +3,38 @@
 load("assert.star", "assert")
 load("proto.star", "proto")
 
-schema = proto.file("google/protobuf/descriptor.proto")
+schema = proto.file("test.proto")
 
-m = schema.FileDescriptorProto(name = "somename.proto", dependency = ["a", "b", "c"])
+m = schema.Test(
+    string_field="I'm a string!",
+    int32_field=42,
+    repeated_field=["a", "b", "c"],
+    map_field={"a": "A", "b": "B"}
+)
 assert.eq(type(m), "proto.Message")
-assert.eq(m.name, "somename.proto")
-assert.eq(list(m.dependency), ["a", "b", "c"])
-m.dependency = ["d", "e"]
-assert.eq(list(m.dependency), ["d", "e"])
-m.dependency.append("f")
-assert.eq(list(m.dependency), ["d", "e", "f"])
+assert.eq(type(m.repeated_field), "proto.repeated<string>")
+assert.eq(type(m.map_field), "proto.map<string, string>")
+
+assert.eq(m.string_field, "I'm a string!")
+assert.eq(m.int32_field, 42)
+
+assert.eq(list(m.repeated_field), ["a", "b", "c"])
+m.repeated_field = ["d", "e"]
+assert.eq(list(m.repeated_field), ["d", "e"])
+m.repeated_field.append("f")
+assert.eq(list(m.repeated_field), ["d", "e", "f"])
+
+assert.eq(dict(m.map_field), {"a": "A", "b": "B"})
+m.map_field["c"] = "C"
+assert.eq(dict(m.map_field), {"a": "A", "b": "B", "c": "C"})
+
+
+# m = schema.FileDescriptorProto(name = "somename.proto", dependency = ["a", "b", "c"])
+# assert.eq(type(m), "proto.Message")
+# assert.eq(m.name, "somename.proto")
+# assert.eq(list(m.dependency), ["a", "b", "c"])
+# m.dependency = ["d", "e"]
+# assert.eq(list(m.dependency), ["d", "e"])
+# m.dependency.append("f")
+# assert.eq(list(m.dependency), ["d", "e", "f"])
 
diff --git a/starlark/testdata/proto/test.fds b/starlark/testdata/proto/test.fds
new file mode 100644
index 00000000..2d94d85a
--- /dev/null
+++ b/starlark/testdata/proto/test.fds
@@ -0,0 +1,13 @@
+
+ª
+
+test.protogo.starlark.net.testdata"ù

+Test!
+string_field
 
(	RstringField
+int32_field 
(R
+int32Field%
+repeated_field (	R
repeatedFieldI
+	map_field (2,.go.starlark.net.testdata.Test.MapFieldEntryRmapField;
+
MapFieldEntry
+key
 
(	Rkey
+value 
(	Rvalue:8
bproto3
\ No newline at end of file
diff --git a/starlark/testdata/proto/test.proto b/starlark/testdata/proto/test.proto
new file mode 100644
index 00000000..52f49024
--- /dev/null
+++ b/starlark/testdata/proto/test.proto
@@ -0,0 +1,11 @@
+syntax = "proto3";
+
+package go.starlark.net.testdata;
+
+message Test {
+    string string_field = 1;
+    int32 int32_field = 2;
+    repeated string repeated_field = 3;
+    map<string, string> map_field = 4;
+}
+