// Copyright (c) 2025 Probo Inc . // // Permission to use, copy, modify, and/or distribute this software for any // purpose with or without fee is hereby granted, provided that the above // copyright notice and this permission notice appear in all copies. // // THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH // REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY // AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, // INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM // LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR // OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR // PERFORMANCE OF THIS SOFTWARE. package testutil import ( "bytes" "crypto/rand" "encoding/hex" "encoding/json" "fmt" "io" "net/http" "net/http/cookiejar" "testing" "time" "github.com/stretchr/testify/require" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" ) func generateUniqueID() string { randomBytes := make([]byte, 4) rand.Read(randomBytes) return fmt.Sprintf("%d-%s", time.Now().UnixNano(), hex.EncodeToString(randomBytes)) } type TestRole string const ( RoleOwner TestRole = "OWNER" RoleAdmin TestRole = "ADMIN" RoleViewer TestRole = "VIEWER" ) type Client struct { T testing.TB httpClient *http.Client baseURL string role TestRole userID gid.GID organizationID gid.GID } func NewClient(t testing.TB, role TestRole) *Client { t.Helper() jar, err := cookiejar.New(nil) require.NoError(t, err, "cannot create cookie jar") client := &Client{ T: t, baseURL: GetBaseURL(), role: role, httpClient: &http.Client{ Jar: jar, Timeout: 30 * time.Second, }, } client.setupTestUser() return client } func NewClientInOrg(t testing.TB, role TestRole, ownerClient *Client) *Client { t.Helper() jar, err := cookiejar.New(nil) require.NoError(t, err, "cannot create cookie jar") client := &Client{ T: t, baseURL: GetBaseURL(), role: role, organizationID: ownerClient.organizationID, httpClient: &http.Client{ Jar: jar, Timeout: 30 * time.Second, }, } client.setupTestUserInOrg(ownerClient) return client } func (c *Client) setupTestUser() { uniqueID := generateUniqueID() email := fmt.Sprintf("test-%s@e2e.probo.test", uniqueID) password := "TestPassword123!" fullName := fmt.Sprintf("Test User %s", uniqueID) // Sign up c.userID = c.signUp(email, password, fullName) // Create organization (this makes the user an OWNER) orgName := fmt.Sprintf("Test Org %s", uniqueID) c.organizationID = c.createOrganization(orgName) // If the role is not OWNER, we need to adjust the membership if c.role != RoleOwner { c.updateOwnMembershipRole(coredata.MembershipRole(c.role)) } } func (c *Client) setupTestUserInOrg(ownerClient *Client) { uniqueID := generateUniqueID() email := fmt.Sprintf("test-%s@e2e.probo.test", uniqueID) password := "TestPassword123!" fullName := fmt.Sprintf("Test User %s", uniqueID) // Sign up new user c.userID = c.signUp(email, password, fullName) // Owner invites user to organization invitationID := ownerClient.inviteMember(email, fullName, coredata.MembershipRole(c.role)) // New user accepts invitation c.acceptInvitation(invitationID) } func (c *Client) signUp(email, password, fullName string) gid.GID { payload := map[string]string{ "email": email, "password": password, "fullName": fullName, } body, err := json.Marshal(payload) require.NoError(c.T, err, "cannot marshal sign-up payload") req, err := http.NewRequest("POST", c.baseURL+"/connect/register", bytes.NewReader(body)) require.NoError(c.T, err, "cannot create sign-up request") req.Header.Set("Content-Type", "application/json") resp, err := c.httpClient.Do(req) require.NoError(c.T, err, "sign-up request failed") defer resp.Body.Close() respBody, _ := io.ReadAll(resp.Body) require.Equal(c.T, http.StatusOK, resp.StatusCode, "sign-up failed: %s", string(respBody)) var result struct { User struct { ID string `json:"id"` } `json:"user"` } err = json.Unmarshal(respBody, &result) require.NoError(c.T, err, "cannot decode sign-up response") userID, err := gid.ParseGID(result.User.ID) require.NoError(c.T, err, "cannot parse user ID") return userID } func (c *Client) createOrganization(name string) gid.GID { const query = ` mutation($input: CreateOrganizationInput!) { createOrganization(input: $input) { organizationEdge { node { id } } } } ` var result struct { CreateOrganization struct { OrganizationEdge struct { Node struct { ID string `json:"id"` } `json:"node"` } `json:"organizationEdge"` } `json:"createOrganization"` } err := c.Execute(query, map[string]any{ "input": map[string]any{"name": name}, }, &result) require.NoError(c.T, err, "createOrganization mutation failed") orgID, err := gid.ParseGID(result.CreateOrganization.OrganizationEdge.Node.ID) require.NoError(c.T, err, "cannot parse organization ID") return orgID } func (c *Client) updateOwnMembershipRole(role coredata.MembershipRole) { // First get the membership ID const queryMemberships = ` query($id: ID!) { organization(id: $id) { memberships(first: 100) { edges { node { id userId role } } } } } ` var qResult struct { Organization struct { Memberships struct { Edges []struct { Node struct { ID string `json:"id"` UserID string `json:"userId"` Role string `json:"role"` } `json:"node"` } `json:"edges"` } `json:"memberships"` } `json:"organization"` } err := c.Execute(queryMemberships, map[string]any{ "id": c.organizationID.String(), }, &qResult) require.NoError(c.T, err, "cannot query organization memberships") var membershipID string for _, edge := range qResult.Organization.Memberships.Edges { if edge.Node.UserID == c.userID.String() { membershipID = edge.Node.ID break } } require.NotEmpty(c.T, membershipID, "membership not found for user") // Update the role const updateQuery = ` mutation($input: UpdateMembershipInput!) { updateMembership(input: $input) { membership { id role } } } ` err = c.Execute(updateQuery, map[string]any{ "input": map[string]any{ "organizationId": c.organizationID.String(), "memberId": membershipID, "role": string(role), }, }, nil) require.NoError(c.T, err, "updateMembership mutation failed") } func (c *Client) inviteMember(email, fullName string, role coredata.MembershipRole) gid.GID { const query = ` mutation($input: InviteUserInput!) { inviteUser(input: $input) { invitationEdge { node { id } } } } ` var result struct { InviteUser struct { InvitationEdge struct { Node struct { ID string `json:"id"` } `json:"node"` } `json:"invitationEdge"` } `json:"inviteUser"` } err := c.Execute(query, map[string]any{ "input": map[string]any{ "organizationId": c.organizationID.String(), "email": email, "fullName": fullName, "role": string(role), "createPeople": false, }, }, &result) require.NoError(c.T, err, "inviteUser mutation failed") invitationID, err := gid.ParseGID(result.InviteUser.InvitationEdge.Node.ID) require.NoError(c.T, err, "cannot parse invitation ID") return invitationID } func (c *Client) acceptInvitation(invitationID gid.GID) { const query = ` mutation($input: AcceptInvitationInput!) { acceptInvitation(input: $input) { invitation { id status } } } ` err := c.Execute(query, map[string]any{ "input": map[string]any{ "invitationId": invitationID.String(), }, }, nil) require.NoError(c.T, err, "acceptInvitation mutation failed") } func (c *Client) GetUserID() gid.GID { return c.userID } func (c *Client) GetOrganizationID() gid.GID { return c.organizationID } func (c *Client) GetRole() TestRole { return c.role }