bug fix in Get method of ColoredBalances (iotaledger#1089)

* bug fix in Get method of ColoredBalances

* move test to color_test.go

* import sorting

packages/ledgerstate/color.go

@@ -175,8 +175,11 @@ func ColoredBalancesFromMarshalUtil(marshalUtil *marshalutil.MarshalUtil) (color
 // Get returns the balance of the given Color and a boolean value indicating if the requested Color existed.
 func (c *ColoredBalances) Get(color Color) (uint64, bool) {
 	balance, exists := c.balances.Get(color)
-
-	return balance.(uint64), exists
+	ret, ok := balance.(uint64)
+	if !ok {
+		return 0, false
+	}
+	return ret, exists
 }
 
 // ForEach calls the consumer for each element in the collection and aborts the iteration if the consumer returns false.

packages/ledgerstate/color_test.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/blake2b"
 )
 
 func TestNewColoredBalances(t *testing.T) {
@@ -67,3 +68,21 @@ func TestColoredBalances_String(t *testing.T) {
 
 	assert.Equal(t, "ColoredBalances {\n    IOTA: 100\n    MINT: 120\n}", coloredBalances.String())
 }
+
+func TestNonExistentColor(t *testing.T) {
+	h := blake2b.Sum256([]byte("dummy"))
+	color, _, err := ColorFromBytes(h[:])
+	require.NoError(t, err)
+	m := map[Color]uint64{color: 5}
+
+	cb := NewColoredBalances(m)
+	amount, ok := cb.Get(color)
+	require.EqualValues(t, 5, amount)
+	require.True(t, ok)
+
+	require.NotPanics(t, func() {
+		amount, ok = cb.Get(ColorIOTA)
+	})
+	require.False(t, ok)
+	require.EqualValues(t, 0, amount)
+}