From 5c4407cbaf7b82104f82a98d6565da358dd450b9 Mon Sep 17 00:00:00 2001
From: Michael de Sa <mjdesa@gmail.com>
Date: Thu, 26 Oct 2017 16:02:57 -0700
Subject: [PATCH] Test MeOrganization for setting and changing current org

Signed-off-by: Jared Scheib <jared.scheib@gmail.com>
---
 server/me_test.go | 123 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 123 insertions(+)

diff --git a/server/me_test.go b/server/me_test.go
index 31db60f014..05a9f966b8 100644
--- a/server/me_test.go
+++ b/server/me_test.go
@@ -1,7 +1,9 @@
 package server
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -179,3 +181,124 @@ func TestService_Me(t *testing.T) {
 		}
 	}
 }
+
+func TestService_MeOrganizations(t *testing.T) {
+	type fields struct {
+		UsersStore chronograf.UsersStore
+		Logger     chronograf.Logger
+		UseAuth    bool
+	}
+	type args struct {
+		w          *httptest.ResponseRecorder
+		r          *http.Request
+		orgRequest *meOrganizationRequest
+		auth       mocks.Authenticator
+	}
+	tests := []struct {
+		name            string
+		fields          fields
+		args            args
+		principal       oauth2.Principal
+		wantStatus      int
+		wantContentType string
+		wantBody        string
+	}{
+		{
+			name: "Set the current User's organization",
+			args: args{
+				w: httptest.NewRecorder(),
+				r: httptest.NewRequest("GET", "http://example.com/foo", nil),
+				orgRequest: &meOrganizationRequest{
+					OrganizationID: "1337",
+				},
+				auth: mocks.Authenticator{},
+			},
+			fields: fields{
+				UseAuth: true,
+				Logger:  log.New(log.DebugLevel),
+				UsersStore: &mocks.UsersStore{
+					GetF: func(ctx context.Context, q chronograf.UserQuery) (*chronograf.User, error) {
+						if q.Name == nil || q.Provider == nil || q.Scheme == nil {
+							return nil, fmt.Errorf("Invalid user query: missing Name, Provider, and/or Scheme")
+						}
+						return &chronograf.User{
+							Name:     "me",
+							Provider: "github",
+							Scheme:   "oauth2",
+						}, nil
+					},
+				},
+			},
+			principal: oauth2.Principal{
+				Subject: "me",
+				Issuer:  "github",
+			},
+			wantStatus:      http.StatusOK,
+			wantContentType: "application/json",
+			wantBody:        `{"name":"me","provider":"github","scheme":"oauth2","currentOrganization":"1337","links":{"self":"/chronograf/v1/users/me"}}`,
+		},
+		{
+			name: "Change the current User's organization",
+			args: args{
+				w: httptest.NewRecorder(),
+				r: httptest.NewRequest("GET", "http://example.com/foo", nil),
+				orgRequest: &meOrganizationRequest{
+					OrganizationID: "1337",
+				},
+				auth: mocks.Authenticator{},
+			},
+			fields: fields{
+				UseAuth: true,
+				Logger:  log.New(log.DebugLevel),
+				UsersStore: &mocks.UsersStore{
+					GetF: func(ctx context.Context, q chronograf.UserQuery) (*chronograf.User, error) {
+						if q.Name == nil || q.Provider == nil || q.Scheme == nil {
+							return nil, fmt.Errorf("Invalid user query: missing Name, Provider, and/or Scheme")
+						}
+						return &chronograf.User{
+							Name:     "me",
+							Provider: "github",
+							Scheme:   "oauth2",
+						}, nil
+					},
+				},
+			},
+			principal: oauth2.Principal{
+				Subject:      "me",
+				Issuer:       "github",
+				Organization: "1338",
+			},
+			wantStatus:      http.StatusOK,
+			wantContentType: "application/json",
+			wantBody:        `{"name":"me","provider":"github","scheme":"oauth2","currentOrganization":"1337","links":{"self":"/chronograf/v1/users/me"}}`,
+		},
+	}
+	for _, tt := range tests {
+		tt.args.r = tt.args.r.WithContext(context.WithValue(context.Background(), oauth2.PrincipalKey, tt.principal))
+		s := &Service{
+			UsersStore: tt.fields.UsersStore,
+			Logger:     tt.fields.Logger,
+			UseAuth:    tt.fields.UseAuth,
+		}
+
+		buf, _ := json.Marshal(tt.args.orgRequest)
+		tt.args.r.Body = ioutil.NopCloser(bytes.NewReader(buf))
+		tt.args.auth.Principal = tt.principal
+
+		s.MeOrganization(&tt.args.auth)(tt.args.w, tt.args.r)
+
+		resp := tt.args.w.Result()
+		content := resp.Header.Get("Content-Type")
+		body, _ := ioutil.ReadAll(resp.Body)
+
+		if resp.StatusCode != tt.wantStatus {
+			t.Errorf("%q. Me() = %v, want %v", tt.name, resp.StatusCode, tt.wantStatus)
+		}
+		if tt.wantContentType != "" && content != tt.wantContentType {
+			t.Errorf("%q. Me() = %v, want %v", tt.name, content, tt.wantContentType)
+		}
+		if eq, err := jsonEqual(tt.wantBody, string(body)); err != nil || !eq {
+			t.Errorf("%q. Me() = \n***%v***\n,\nwant\n***%v***", tt.name, string(body), tt.wantBody)
+		}
+	}
+}