Add special union type handling

Author: rocktavious

relationship_test.go

@@ -147,15 +147,17 @@ func TestRelationshipDefinitionDelete(t *testing.T) {
 func TestGetRelationship(t *testing.T) {
 	// Arrange
 	testRequest := autopilot.NewTestRequest(
-		`query RelationshipGet($input:ID!){account{relationship(id: $input){id,source{... on Domain{id,aliases},... on InfrastructureResource{id,aliases,name},... on Service{id,aliases},... on System{id,aliases},... on Team{alias,id}},target{... on Domain{id,aliases},... on InfrastructureResource{id,aliases,name},... on Service{id,aliases},... on System{id,aliases},... on Team{alias,id}},type}}}`,
+		`query RelationshipGet($input:ID!){account{relationship(id: $input){id,source{__typename,... on Domain{id,aliases},... on InfrastructureResource{id,aliases,name},... on Service{id,aliases},... on System{id,aliases},... on Team{alias,id}},target{__typename,... on Domain{id,aliases},... on InfrastructureResource{id,aliases,name},... on Service{id,aliases},... on System{id,aliases},... on Team{alias,id}},type}}}`,
 		`{"input": "{{ template "id1_string" }}" }`,
 		`{"data": {"account": {"relationship": {
 			"id": "{{ template "id1_string" }}",
 			"source": {
+				"__typename": "Domain",
 				"id": "{{ template "id2_string" }}",
 				"aliases": ["source-alias"]
 			},
 			"target": {
+				"__typename": "System",
 				"id": "{{ template "id3_string" }}",
 				"aliases": ["target-alias"]
 			},
@@ -170,6 +172,8 @@ func TestGetRelationship(t *testing.T) {
 	autopilot.Ok(t, err)
 	autopilot.Equals(t, id1, result.Id)
 	autopilot.Equals(t, id2, result.Source.Domain.Id)
-	autopilot.Equals(t, id3, result.Target.Domain.Id)
+	autopilot.Equals(t, ol.ID(""), result.Source.System.Id)
+	autopilot.Equals(t, ol.ID(""), result.Target.Domain.Id)
+	autopilot.Equals(t, id3, result.Target.System.Id)
 	autopilot.Equals(t, ol.RelationshipTypeEnumBelongsTo, result.Type)
 }

union.go

@@ -0,0 +1,65 @@
+package opslevel
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+// TypeIdentifier represents a GraphQL __typename that can be used to identify the concrete type of an union response
+type TypeIdentifier string
+
+func (r *RelationshipResource) MarshalJSON() ([]byte, error) {
+	switch r.Type {
+	case "Domain":
+		return json.Marshal(r.Domain)
+	case "Infrastructure":
+		return json.Marshal(r.InfrastructureResource)
+	case "Service":
+		return json.Marshal(r.Service)
+	case "System":
+		return json.Marshal(r.System)
+	case "Team":
+		return json.Marshal(r.Team)
+	default:
+		return nil, fmt.Errorf("unknown type: %s", r.Type)
+	}
+}
+
+// UnmarshalJSON implements custom JSON unmarshaling for RelationshipResource to support the union type
+func (r *RelationshipResource) UnmarshalJSON(data []byte) error {
+	// First unmarshal to get the type
+	var tmp struct {
+		Type TypeIdentifier `json:"__typename"`
+	}
+	if err := json.Unmarshal(data, &tmp); err != nil {
+		return fmt.Errorf("failed to unmarshal type: %w", err)
+	}
+	r.Type = tmp.Type
+
+	// Then unmarshal the specific type
+	switch string(r.Type) {
+	case "Domain":
+		if err := json.Unmarshal(data, &r.Domain); err != nil {
+			return fmt.Errorf("failed to unmarshal Domain: %w", err)
+		}
+	case "Infrastructure":
+		if err := json.Unmarshal(data, &r.InfrastructureResource); err != nil {
+			return fmt.Errorf("failed to unmarshal Infrastructure Resource: %w", err)
+		}
+	case "Service":
+		if err := json.Unmarshal(data, &r.Service); err != nil {
+			return fmt.Errorf("failed to unmarshal Service: %w", err)
+		}
+	case "System":
+		if err := json.Unmarshal(data, &r.System); err != nil {
+			return fmt.Errorf("failed to unmarshal System: %w", err)
+		}
+	case "Team":
+		if err := json.Unmarshal(data, &r.Team); err != nil {
+			return fmt.Errorf("failed to unmarshal Team: %w", err)
+		}
+	default:
+		return fmt.Errorf("unknown resource type: %s", r.Type)
+	}
+	return nil
+}

union_type.go

@@ -0,0 +1,56 @@
+package opslevel_test
+
+import (
+	"encoding/json"
+	ol "github.com/opslevel/opslevel-go/v2025"
+	"github.com/rocktavious/autopilot/v2023"
+	"testing"
+)
+
+func TestRelationshipResourceUnmarshalJSON(t *testing.T) {
+	// Arrange
+	testCases := []struct {
+		name     string
+		input    string
+		expected ol.RelationshipResource
+	}{
+		{
+			name:  "Domain resource",
+			input: `{"__typename": "Domain", "id": "Z2lkOi8vb3BzbGV2ZWwvRG9tYWluLzE"}`,
+			expected: ol.RelationshipResource{
+				Type: ol.TypeIdentifier("Domain"),
+				Domain: ol.DomainId{
+					Id: "Z2lkOi8vb3BzbGV2ZWwvRG9tYWluLzE",
+				},
+			},
+		},
+		{
+			name:  "Service resource",
+			input: `{"__typename": "Service", "id": "Z2lkOi8vb3BzbGV2ZWwvU2VydmljZS8x"}`,
+			expected: ol.RelationshipResource{
+				Type: ol.TypeIdentifier("Service"),
+				Service: ol.ServiceId{
+					Id: ol.ID("Z2lkOi8vb3BzbGV2ZWwvU2VydmljZS8x"),
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Act
+			var result ol.RelationshipResource
+			err := json.Unmarshal([]byte(tc.input), &result)
+
+			// Assert
+			autopilot.Ok(t, err)
+			autopilot.Equals(t, tc.expected.Type, result.Type)
+			switch tc.expected.Type {
+			case "Domain":
+				autopilot.Equals(t, tc.expected.Domain, result.Domain)
+			case "Service":
+				autopilot.Equals(t, tc.expected.Service.Id, result.Service.Id)
+			}
+		})
+	}
+}