| // Copyright 2020 Google LLC |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| // |
| |
| package pbeam |
| |
| import ( |
| "reflect" |
| "testing" |
| |
| "github.com/google/differential-privacy/privacy-on-beam/v2/pbeam/testutils" |
| testpb "github.com/google/differential-privacy/privacy-on-beam/v2/testdata" |
| "github.com/apache/beam/sdks/v2/go/pkg/beam" |
| "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert" |
| "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest" |
| "github.com/google/go-cmp/cmp" |
| "google.golang.org/protobuf/proto" |
| ) |
| |
| func init() { |
| beam.RegisterType(reflect.TypeOf((*testpb.TestAnon)(nil))) |
| beam.RegisterType(reflect.TypeOf(protoPair{})) |
| } |
| |
| type protoPair struct { |
| Key string |
| Pb *testpb.TestAnon |
| } |
| |
| func kvToProtoPair(key string, pb *testpb.TestAnon) protoPair { |
| return protoPair{key, pb} |
| } |
| |
| func TestMakePrivate(t *testing.T) { |
| values := []testutils.PairII{ |
| {17, 42}, |
| {99, 0}, |
| } |
| p, s, col := ptest.CreateList(values) |
| colKV := beam.ParDo(s, testutils.PairToKV, col) |
| |
| // pcol should contain 17→42 and 99→0. |
| pcol := MakePrivate(s, colKV, NewPrivacySpec(1, 1e-10)) |
| got := beam.ParDo(s, testutils.KVToPair, pcol.col) |
| passert.Equals(s, got, col) |
| if err := ptest.Run(p); err != nil { |
| t.Errorf("MakePrivate(%v) = %v, expected %v: %v", col, got, col, err) |
| } |
| } |
| |
| type SimpleStruct struct { |
| String string |
| Int int |
| } |
| |
| type ComplexStruct struct { |
| String string |
| Int int |
| StringPointer *string |
| StringSlice []string |
| SubStruct *SimpleStruct |
| SubStructSlice []SimpleStruct |
| } |
| |
| type structPair struct { |
| Key string |
| Value ComplexStruct |
| } |
| |
| func kvToStructPair(key string, value ComplexStruct) structPair { |
| return structPair{Key: key, Value: value} |
| } |
| |
| type RecursiveStruct struct { |
| String string |
| Int int |
| SubStruct ComplexStruct |
| RecursiveStruct *RecursiveStruct |
| } |
| |
| func TestMakePrivateFromStruct(t *testing.T) { |
| fortyTwo := "42" |
| seventeen := "17" |
| for _, tc := range []struct { |
| desc string |
| idFieldPath string |
| values []ComplexStruct |
| want []structPair |
| }{ |
| {"top level string id field", |
| "String", |
| []ComplexStruct{ |
| {String: "42", Int: 42}, |
| {String: "17", Int: 17}}, |
| []structPair{ |
| {Key: "\"42\"", Value: ComplexStruct{String: "42", Int: 42}}, |
| {Key: "\"17\"", Value: ComplexStruct{String: "17", Int: 17}}}, |
| }, |
| {"top level string pointer id field", |
| "StringPointer", |
| []ComplexStruct{ |
| {StringPointer: &fortyTwo, Int: 42}, |
| {StringPointer: &seventeen, Int: 17}}, |
| []structPair{ |
| {Key: "\"42\"", Value: ComplexStruct{StringPointer: &fortyTwo, Int: 42}}, |
| {Key: "\"17\"", Value: ComplexStruct{StringPointer: &seventeen, Int: 17}}}, |
| }, |
| {"bottom level string id field", |
| "SubStruct.String", |
| []ComplexStruct{ |
| {SubStruct: &SimpleStruct{String: "42"}, Int: 42}, |
| {SubStruct: &SimpleStruct{String: "17"}, Int: 17}}, |
| []structPair{ |
| {Key: "\"42\"", Value: ComplexStruct{SubStruct: &SimpleStruct{String: "42"}, Int: 42}}, |
| {Key: "\"17\"", Value: ComplexStruct{SubStruct: &SimpleStruct{String: "17"}, Int: 17}}}, |
| }, |
| } { |
| p, s, col, want := ptest.CreateList2(tc.values, tc.want) |
| |
| pcol := MakePrivateFromStruct(s, col, NewPrivacySpec(1, 1e-10), tc.idFieldPath) |
| got := beam.ParDo(s, kvToStructPair, pcol.col) |
| passert.Equals(s, got, want) |
| if err := ptest.Run(p); err != nil { |
| t.Errorf("MakePrivateFromStruct output does not match input values with %s. got %v, expected %v: %v", tc.desc, got, want, err) |
| } |
| } |
| } |
| |
| // Tests the GetIDField method in extractStructFieldFn. |
| func TestGetIDField(t *testing.T) { |
| eight := "8" |
| val := RecursiveStruct{ |
| String: "0", |
| Int: 0, |
| SubStruct: ComplexStruct{ |
| String: "1", |
| Int: 1, |
| StringSlice: []string{"2", "3", "4"}, |
| SubStruct: &SimpleStruct{String: "5", Int: 5}, |
| SubStructSlice: []SimpleStruct{ |
| SimpleStruct{String: "6", Int: 6}, |
| SimpleStruct{String: "7", Int: 7}, |
| }, |
| StringPointer: &eight}, |
| RecursiveStruct: &RecursiveStruct{String: "9", Int: 9}, |
| } |
| for _, tc := range []struct { |
| idFieldPath string |
| want interface{} |
| wantErr bool |
| }{ |
| {"String", "0", false}, |
| {"Int", 0, false}, |
| {"SubStruct", nil, true}, |
| {"SubStruct.String", "1", false}, |
| {"SubStruct.Int", 1, false}, |
| {"SubStruct.StringSlice", nil, true}, |
| {"SubStruct.SubStruct", nil, true}, |
| {"SubStruct.SubStruct.String", "5", false}, |
| {"SubStruct.SubStruct.Int", 5, false}, |
| {"SubStruct.SubStructSlice", nil, true}, |
| {"SubStruct.SubStructSlice.String", nil, true}, |
| {"SubStruct.SubStructSlice.Int", nil, true}, |
| {"SubStruct.StringPointer", "8", false}, |
| {"RecursiveStruct", nil, true}, |
| {"RecursiveStruct.String", "9", false}, |
| {"RecursiveStruct.Int", 9, false}, |
| {"RecursiveStruct.RecursiveStruct", nil, true}, |
| {"RecursiveStruct.RecursiveStruct.String", "", false}, |
| {"RecursiveStruct.RecursiveStruct.Int", 0, false}, |
| {"nonexistent", nil, true}, |
| } { |
| ext := extractStructFieldFn{IDFieldPath: tc.idFieldPath} |
| got, err := ext.getIDField(val) |
| if (err != nil) != tc.wantErr { |
| t.Errorf("GetIDField with idFieldPath=%s: got error %v, wantErr=%t.", tc.idFieldPath, err, tc.wantErr) |
| } |
| if !cmp.Equal(got, tc.want) { |
| t.Errorf("GetIDField with idFieldPath=%s: retrieved field %v, wanted=%v.", tc.idFieldPath, got, tc.want) |
| } |
| } |
| } |
| |
| func TestMakePrivateFromProto(t *testing.T) { |
| values := []*testpb.TestAnon{ |
| &testpb.TestAnon{Foo: proto.Int64(42), Bar: proto.String("fourty-two")}, |
| &testpb.TestAnon{Foo: proto.Int64(17), Bar: proto.String("seventeen")}, |
| &testpb.TestAnon{Bar: proto.String("zero")}, |
| } |
| result := []protoPair{ |
| {"42", &testpb.TestAnon{Foo: proto.Int64(42), Bar: proto.String("fourty-two")}}, |
| {"17", &testpb.TestAnon{Foo: proto.Int64(17), Bar: proto.String("seventeen")}}, |
| {"0", &testpb.TestAnon{Bar: proto.String("zero")}}, |
| } |
| p, s, col, want := ptest.CreateList2(values, result) |
| |
| pcol := MakePrivateFromProto(s, col, NewPrivacySpec(1, 1e-10), "foo") |
| got := beam.ParDo(s, kvToProtoPair, pcol.col) |
| passert.Equals(s, got, want) |
| if err := ptest.Run(p); err != nil { |
| t.Errorf("MakePrivateFromProto(%v) = %v, expected %v: %v", col, got, want, err) |
| } |
| } |
| |
| var ( |
| repeat = []string{"bar", "baz"} |
| subrepeat = []*testpb.TestComplex_Submessage{ |
| &testpb.TestComplex_Submessage{ |
| Simple: proto.String("oob"), |
| Repeat: repeat, |
| }, |
| &testpb.TestComplex_Submessage{ |
| Simple: proto.String("obo"), |
| Repeat: repeat, |
| }, |
| } |
| complexMsg = &testpb.TestComplex{ |
| Simple: proto.String("foo"), |
| Repeat: repeat, |
| Sub: &testpb.TestComplex_Submessage{ |
| Simple: proto.String("boo"), |
| Repeat: repeat, |
| }, |
| Subrepeat: subrepeat, |
| } |
| withoutSimple = &testpb.TestComplex{ |
| Repeat: repeat, |
| Sub: &testpb.TestComplex_Submessage{ |
| Simple: proto.String("boo"), |
| Repeat: repeat, |
| }, |
| Subrepeat: subrepeat, |
| } |
| withoutSubSimple = &testpb.TestComplex{ |
| Simple: proto.String("foo"), |
| Repeat: repeat, |
| Sub: &testpb.TestComplex_Submessage{ |
| Repeat: repeat, |
| }, |
| Subrepeat: subrepeat, |
| } |
| ) |
| |
| // Tests the extraction logic in extractProtoFieldFn. |
| func TestExtractProtoField(t *testing.T) { |
| for _, tc := range []struct { |
| idFieldPath string |
| wantField string |
| wantMsg *testpb.TestComplex |
| ok bool |
| }{ |
| {"simple", "foo", complexMsg, true}, |
| {"empty", "", complexMsg, true}, |
| {"sub.simple", "boo", complexMsg, true}, |
| {"repeat", "", nil, false}, |
| {"sub.repeat", "", nil, false}, |
| {"subrepeat.simple", "", nil, false}, |
| {"subrepeat.repeat", "", nil, false}, |
| {"nonexistent", "", nil, false}, |
| } { |
| ext := &extractProtoFieldFn{ |
| IDFieldPath: tc.idFieldPath, |
| desc: (&testpb.TestComplex{}).ProtoReflect().Descriptor(), |
| } |
| clone := &testpb.TestComplex{} |
| proto.Merge(clone, complexMsg) |
| gotField, err := ext.extractField(clone.ProtoReflect()) |
| if (err == nil) != tc.ok { |
| t.Errorf("extractField with IDFieldPath=%s: got error %v, want ok=%t.", tc.idFieldPath, err, tc.ok) |
| } |
| if err == nil { |
| gotField := gotField.(string) |
| if gotField != tc.wantField { |
| t.Errorf("extractField with IDFieldPath=%s: got field %v, want %s", tc.idFieldPath, gotField, tc.wantField) |
| } |
| if !proto.Equal(clone, tc.wantMsg) { |
| t.Errorf("extractField with IDFieldPath=%s: got msg %v, want %v", tc.idFieldPath, clone, tc.wantMsg) |
| } |
| } |
| } |
| } |
| |
| // Tests that we can get get the whole budget and consume it partially afterwards. |
| func TestGetFullBudget(t *testing.T) { |
| spec := NewPrivacySpec(2, 2e-10) |
| eps, del, err := spec.getBudget(0, 0) |
| if err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| if eps != 2.0 || del != 2e-10 { |
| t.Errorf("Trying to get the whole budget: Got (epsilon,delta)=(%f,%e), expected=(%f,%e)", eps, del, 2.0, 2e-10) |
| } |
| |
| // Split the budget and consume it in two calls. |
| eps, del, err = spec.consumeBudget(1, 1e-10) |
| if err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| if eps != 1.0 || del != 1e-10 { |
| t.Errorf("Trying to consume the budget after getBudget call: Got (epsilon,delta)=(%f,%e), expected=(%f,%e)", eps, del, 1.0, 1e-10) |
| } |
| eps, del, err = spec.consumeBudget(1, 1e-10) |
| if err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| if eps != 1.0 || del != 1e-10 { |
| t.Errorf("Trying to consume the budget after getBudget call: Got (epsilon,delta)=(%f,%e), expected=(%f,%e)", eps, del, 1.0, 1e-10) |
| } |
| } |
| |
| // Tests that we can get and consume the budget partially. |
| func TestGetPartialBudget(t *testing.T) { |
| spec := NewPrivacySpec(2, 2e-10) |
| eps, del, err := spec.getBudget(1, 1e-10) |
| if err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| if eps != 1.0 || del != 1e-10 { |
| t.Errorf("Trying to get first half of the budget: Got (epsilon,delta)=(%f,%e), expected=(%f,%e)", eps, del, 1.0, 1e-10) |
| } |
| |
| eps, del, err = spec.consumeBudget(1, 1e-10) |
| if err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| if eps != 1.0 || del != 1e-10 { |
| t.Errorf("Trying to consume second half of the budget after getBudget call: Got (epsilon,delta)=(%f,%e), expected=(%f,%e)", eps, del, 1.0, 1e-10) |
| } |
| |
| eps, del, err = spec.getBudget(1, 1e-10) |
| if err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| if eps != 1.0 || del != 1e-10 { |
| t.Errorf("Trying to get second half of the budget: Got (epsilon,delta)=(%f,%e), expected=(%f,%e)", eps, del, 1.0, 1e-10) |
| } |
| |
| eps, del, err = spec.consumeBudget(1, 1e-10) |
| if err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| if eps != 1.0 || del != 1e-10 { |
| t.Errorf("Trying to consume second half the budget after getBudget call: Got (epsilon,delta)=(%f,%e), expected=(%f,%e)", eps, del, 1.0, 1e-10) |
| } |
| } |
| |
| // Tests that we can consume all the budget at once. |
| func TestBudgetFullyConsumed(t *testing.T) { |
| values := []testutils.PairII{ |
| {1, 1}, |
| {2, 2}, |
| } |
| p, s, col := ptest.CreateList(values) |
| colKV := beam.ParDo(s, testutils.PairToKV, col) |
| spec := NewPrivacySpec(1, 1e-30) |
| pcol := MakePrivate(s, colKV, spec) |
| got := Count(s, pcol, CountParams{MaxValue: 1, MaxPartitionsContributed: 1, NoiseKind: LaplaceNoise{}}) |
| passert.Empty(s, got) |
| if err := ptest.Run(p); err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| // Try consuming 1% of the initial budget. |
| if eps, del, err := spec.consumeBudget(0.01, 1e-32); err == nil { |
| t.Errorf("expected spec to be out of budget, but could consume (%f,%e) without any error", eps, del) |
| } |
| } |
| |
| // Tests that two distinct budgets can be independently consumed. |
| func TestTwoDistinctBudgets(t *testing.T) { |
| values := []testutils.PairII{ |
| {1, 1}, |
| {2, 2}, |
| } |
| p, s, col := ptest.CreateList(values) |
| colKV := beam.ParDo(s, testutils.PairToKV, col) |
| spec1 := NewPrivacySpec(1, 1e-30) |
| spec2 := NewPrivacySpec(1, 1e-30) |
| pcol1 := MakePrivate(s, colKV, spec1) |
| pcol2 := MakePrivate(s, colKV, spec2) |
| got1 := Count(s, pcol1, CountParams{MaxValue: 1, MaxPartitionsContributed: 1, NoiseKind: LaplaceNoise{}}) |
| got2 := Count(s, pcol2, CountParams{MaxValue: 1, MaxPartitionsContributed: 1, NoiseKind: LaplaceNoise{}}) |
| passert.Empty(s, got1) |
| passert.Empty(s, got2) |
| if err := ptest.Run(p); err != nil { |
| t.Errorf("expected no error but got error: %v", err) |
| } |
| // Try consuming 1% of the initial budget independently for ε and δ. |
| if eps, del, err := spec1.consumeBudget(0, 1e-32); err == nil { |
| t.Errorf("expected spec1 to be out of budget, but could consume (%f,%e) without any error", eps, del) |
| } |
| if eps, del, err := spec2.consumeBudget(0.01, 0); err == nil { |
| t.Errorf("expected spec2 to be out of budget, but could consume (%f,%e) without any error", eps, del) |
| } |
| } |
| |
| // Test for rounding errors during budget allocation. Dividing the overall |
| // epsilon by 3 leads to rounding errors in this test case. Should run without |
| // any errors. |
| func TestBudgetRounding(t *testing.T) { |
| for numAggregations := 1; numAggregations <= 10; numAggregations++ { |
| values := []testutils.PairII{ |
| {1, 1}, |
| {2, 2}, |
| } |
| p, s, col := ptest.CreateList(values) |
| colKV := beam.ParDo(s, testutils.PairToKV, col) |
| spec := NewPrivacySpec(1, 1e-30) |
| pcol := MakePrivate(s, colKV, spec) |
| epsPerAggregation := 1. / float64(numAggregations) |
| delPerAggregation := 1e-30 / float64(numAggregations) |
| for i := 0; i < numAggregations; i++ { |
| DistinctPrivacyID(s, pcol, DistinctPrivacyIDParams{Epsilon: epsPerAggregation, Delta: delPerAggregation, MaxPartitionsContributed: 1, NoiseKind: LaplaceNoise{}}) |
| } |
| if err := ptest.Run(p); err != nil { |
| t.Errorf("with %d aggregations, expected no error but got error: %v", numAggregations, err) |
| } |
| // Now, the budget should be really empty. |
| if eps, del, err := spec.consumeBudget(1e-15, 1e-40); err == nil { |
| t.Errorf("with %d aggregations, expected spec to be out of budget, but could consume (%f,%e) without any error", numAggregations, eps, del) |
| } |
| } |
| } |
