From 93741082889c73821dbc682f16261d365651229f Mon Sep 17 00:00:00 2001
From: NiseVoid <nisevoid@gmail.com>
Date: Tue, 22 May 2018 12:07:18 +0200
Subject: [PATCH] Fix incorrect attribute conversion

---
 server/main.go      | 15 ++++++++++-----
 server/main_test.go | 20 ++++++++++++++++++++
 2 files changed, 30 insertions(+), 5 deletions(-)
 create mode 100644 server/main_test.go

diff --git a/server/main.go b/server/main.go
index d2291f7..3e90d57 100644
--- a/server/main.go
+++ b/server/main.go
@@ -43,11 +43,7 @@ func pmax(c echo.Context) error {
 	if len(attr) < 2 {
 		return c.NoContent(400)
 	}
-	if len(attr) == 2 {
-		attr = strings.ToUpper(attr[:2])
-	} else {
-		attr = strings.ToUpper(attr[:2]) + strings.ToLower(attr[2:])
-	}
+	attr = convertAttribute(attr)
 
 	date, err := time.Parse(time.RFC3339, c.Param(`date`))
 	if err != nil {
@@ -61,3 +57,12 @@ func pmax(c echo.Context) error {
 
 	return c.JSON(200, data)
 }
+
+func convertAttribute(attr string) string {
+	attr = strings.ToUpper(attr[:1]) + strings.ToLower(attr[1:])
+
+	attr = strings.Replace(attr, `gem`, `Gem`, -1)
+	attr = strings.Replace(attr, `max`, `Max`, -1)
+
+	return attr
+}
diff --git a/server/main_test.go b/server/main_test.go
new file mode 100644
index 0000000..507655a
--- /dev/null
+++ b/server/main_test.go
@@ -0,0 +1,20 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConvertAttribute(t *testing.T) {
+	var data = map[string]string{
+		`cgem`: `CGem`,
+		`ep`:   `Ep`,
+		`PMAX`: `PMax`,
+		`Test`: `Test`,
+	}
+
+	for in, expected := range data {
+		assert.Equal(t, expected, convertAttribute(in))
+	}
+}