Refactor get_tree_public_key test helper methods

cmd/get_tree_public_key/main_test.go

@@ -16,22 +16,18 @@ package main
 
 import (
 	"context"
-	"flag"
 	"fmt"
 	"strings"
 	"testing"
-	"time"
 
-	"github.com/golang/protobuf/ptypes"
 	"github.com/google/trillian"
 	"github.com/google/trillian/client"
-	ktestonly "github.com/google/trillian/crypto/keys/testonly"
-	"github.com/google/trillian/crypto/keyspb"
-	"github.com/google/trillian/crypto/sigpb"
-	"github.com/google/trillian/testonly"
 	"github.com/google/trillian/testonly/integration"
+	"github.com/google/trillian/testonly/setup"
 	"github.com/google/trillian/util/flagsaver"
 	"google.golang.org/grpc"
+
+	stestonly "github.com/google/trillian/storage/testonly"
 )
 
 func TestGetTreePublicKey(t *testing.T) {
@@ -46,67 +42,26 @@ func TestGetTreePublicKey(t *testing.T) {
 	defer logEnv.Close()
 
 	// Create a new Trillian log
-	log := createLog(t, logEnv)
+	log, err := client.CreateAndInitTree(context.Background(), &trillian.CreateTreeRequest{
+		Tree: stestonly.LogTree,
+	}, logEnv.Admin, nil, logEnv.Log)
+	if err != nil {
+		t.Errorf("Failed to create test tree: %v", err)
+	}
 
 	// Set the flags.
 	defer flagsaver.Save().Restore()
-	setFlag(t, "admin_server", logEnv.Address)
-	setFlag(t, "log_id", fmt.Sprint(log.TreeId))
+	setup.SetFlag(t, "admin_server", logEnv.Address)
+	setup.SetFlag(t, "log_id", fmt.Sprint(log.TreeId))
 
 	publicKeyPEM, err := getPublicKeyPEM()
 	if err != nil {
 		t.Errorf("Got an unexpected error: %v", err)
 	}
 
 	// Check that the returned public key PEM is the one we expected.
-	expectedPublicKeyPEM := strings.TrimSpace(testonly.DemoPublicKey)
+	expectedPublicKeyPEM := strings.TrimSpace(stestonly.PublicKeyPEM)
 	if strings.TrimSpace(publicKeyPEM) != expectedPublicKeyPEM {
 		t.Errorf("Expected the public key PEM to equal:\n%s\nInstead got:\n%s", expectedPublicKeyPEM, publicKeyPEM)
 	}
 }
-
-func setFlag(t *testing.T, name, value string) {
-	t.Helper()
-
-	if err := flag.Set(name, value); err != nil {
-		t.Errorf("failed to set the -%s flag: %v", name, err)
-	}
-}
-
-func createLog(t *testing.T, logEnv *integration.LogEnv) *trillian.Tree {
-	t.Helper()
-
-	ctx := context.Background()
-
-	privateKey, err := ptypes.MarshalAny(&keyspb.PrivateKey{
-		Der: ktestonly.MustMarshalPrivatePEMToDER(testonly.DemoPrivateKey, testonly.DemoPrivateKeyPass),
-	})
-	if err != nil {
-		t.Errorf("failed to marshal private key as an any.Any proto: %v", err)
-	}
-
-	tree, err := client.CreateAndInitTree(ctx, &trillian.CreateTreeRequest{
-		Tree: &trillian.Tree{
-			DisplayName:        "Test Log",
-			Description:        "This is a test log.",
-			TreeType:           trillian.TreeType_LOG,
-			TreeState:          trillian.TreeState_ACTIVE,
-			HashStrategy:       trillian.HashStrategy_RFC6962_SHA256,
-			SignatureAlgorithm: sigpb.DigitallySigned_ECDSA,
-			HashAlgorithm:      sigpb.DigitallySigned_SHA256,
-			MaxRootDuration:    ptypes.DurationProto(0 * time.Millisecond),
-
-			// Explicitly set the public and private keys for the new tree.
-			PrivateKey: privateKey,
-			PublicKey: &keyspb.PublicKey{
-				Der: ktestonly.MustMarshalPublicPEMToDER(testonly.DemoPublicKey),
-			},
-		},
-	}, logEnv.Admin, nil, logEnv.Log)
-
-	if err != nil {
-		t.Errorf("failed to create a new log: %v", err)
-	}
-
-	return tree
-}

storage/testonly/admin_storage_tester.go

@@ -58,7 +58,8 @@ Xy3zzHFwlFwjE8L1NCngJAFbu3zFf4IbBOCsz6Fa790utVNdulZncNCl2FMK3U2T
 sdoiTW8ymO+qgwcNrqvPVmjFRBtkN0Pn5lgbWhN/aK3TlS9IYJ/EShbMUzjgVzie
 S9+/31whWcH/FLeLJx4cBzvhgCtfquwA+s5ojeLYYsk=
 -----END EC PRIVATE KEY-----`
-	publicKeyPEM = `
+	// PublicKeyPEM is the public key for: LogTree and PreorderedLogTree
+	PublicKeyPEM = `
 -----BEGIN PUBLIC KEY-----
 MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEywnWicNEQ8bn3GXcGpA+tiU4VL70
 Ws9xezgQPrg96YGsFrF6KYG68iqyHDlQ+4FWuKfGKXHn3ooVtB/pfawb5Q==
@@ -101,7 +102,7 @@ var (
 			Der: ktestonly.MustMarshalPrivatePEMToDER(privateKeyPEM, privateKeyPass),
 		}),
 		PublicKey: &keyspb.PublicKey{
-			Der: ktestonly.MustMarshalPublicPEMToDER(publicKeyPEM),
+			Der: ktestonly.MustMarshalPublicPEMToDER(PublicKeyPEM),
 		},
 		MaxRootDuration: ptypes.DurationProto(0 * time.Millisecond),
 	}
@@ -127,7 +128,7 @@ var (
 			Der: ktestonly.MustMarshalPrivatePEMToDER(privateKeyPEM, privateKeyPass),
 		}),
 		PublicKey: &keyspb.PublicKey{
-			Der: ktestonly.MustMarshalPublicPEMToDER(publicKeyPEM),
+			Der: ktestonly.MustMarshalPublicPEMToDER(PublicKeyPEM),
 		},
 		MaxRootDuration: ptypes.DurationProto(0 * time.Millisecond),
 	}

testonly/setup/flags.go

@@ -0,0 +1,29 @@
+// Copyright 2018 Google Inc. All Rights Reserved.
+//
+// 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 setup
+
+import (
+	"flag"
+	"testing"
+)
+
+// SetFlag updates a flag value, failing the test if something goes wrong.
+func SetFlag(t *testing.T, name, value string) {
+	t.Helper()
+
+	if err := flag.Set(name, value); err != nil {
+		t.Errorf("failed to set the -%s flag: %v", name, err)
+	}
+}