Add ability to mock out the azure query result in the unit test suite

fn.go

@@ -20,10 +20,17 @@ import (
 // TargetXRStatusField is the target field to write the query result to
 const TargetXRStatusField = "status.azResourceGraphQueryResult"
 
+// AzureQueryInterface defines the methods required for querying Azure resources.
+type AzureQueryInterface interface {
+	azQuery(ctx context.Context, azureCreds map[string]string, in *v1beta1.Input) (armresourcegraph.ClientResourcesResponse, error)
+}
+
 // Function returns whatever response you ask it to.
 type Function struct {
 	fnv1.UnimplementedFunctionRunnerServiceServer
 
+	azureQuery AzureQueryInterface
+
 	log logging.Logger
 }
 
@@ -54,60 +61,19 @@ func (f *Function) RunFunction(ctx context.Context, req *fnv1.RunFunctionRequest
 		return rsp, nil
 	}
 
-	var azureCreds map[string]string
-	rawCreds := req.GetCredentials()
-
-	if credsData, ok := rawCreds["azure-creds"]; ok {
-		credsData := credsData.GetCredentialData().GetData()
-		if credsJSON, ok := credsData["credentials"]; ok {
-			err := json.Unmarshal(credsJSON, &azureCreds)
-			if err != nil {
-				response.Fatal(rsp, errors.Wrap(err, "cannot parse json credentials"))
-				return rsp, nil
-			}
-		}
-	} else {
-		response.Fatal(rsp, errors.New("failed to get azure-creds credentials"))
-		return rsp, nil
-	}
-
-	tenantID := azureCreds["tenantId"]
-	clientID := azureCreds["clientId"]
-	clientSecret := azureCreds["clientSecret"]
-	subscriptionID := azureCreds["subscriptionId"]
-
-	// To configure DefaultAzureCredential to authenticate a user-assigned managed identity,
-	// set the environment variable AZURE_CLIENT_ID to the identity's client ID.
-
-	cred, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, nil)
-	if err != nil {
-		response.Fatal(rsp, errors.Wrap(err, "failed to obtain a credentials"))
-		return rsp, nil
-	}
-
-	// Create and authorize a ResourceGraph client
-	client, err := armresourcegraph.NewClient(cred, nil)
+	azureCreds, err := getCreds(req)
 	if err != nil {
-		response.Fatal(rsp, errors.Wrap(err, "failed to create client"))
+		response.Fatal(rsp, err)
 		return rsp, nil
 	}
 
-	queryRequest := armresourcegraph.QueryRequest{
-		Query: to.Ptr(in.Query),
-	}
-
-	if len(subscriptionID) > 0 {
-		queryRequest.Subscriptions = []*string{to.Ptr(subscriptionID)}
-	}
-
-	if len(in.ManagementGroups) > 0 {
-		queryRequest.ManagementGroups = in.ManagementGroups
+	if f.azureQuery == nil {
+		f.azureQuery = &AzureQuery{}
 	}
 
-	// Create the query request, Run the query and get the results.
-	results, err := client.Resources(ctx, queryRequest, nil)
+	results, err := f.azureQuery.azQuery(ctx, azureCreds, in)
 	if err != nil {
-		response.Fatal(rsp, errors.Wrap(err, "failed to finish the request"))
+		response.Fatal(rsp, err)
 		f.log.Info("FAILURE: ", "failure", fmt.Sprint(err))
 		return rsp, nil
 	}
@@ -163,3 +129,66 @@ func (f *Function) RunFunction(ctx context.Context, req *fnv1.RunFunctionRequest
 
 	return rsp, nil
 }
+
+func getCreds(req *fnv1.RunFunctionRequest) (map[string]string, error) {
+	var azureCreds map[string]string
+	rawCreds := req.GetCredentials()
+
+	if credsData, ok := rawCreds["azure-creds"]; ok {
+		credsData := credsData.GetCredentialData().GetData()
+		if credsJSON, ok := credsData["credentials"]; ok {
+			err := json.Unmarshal(credsJSON, &azureCreds)
+			if err != nil {
+				return nil, errors.Wrap(err, "cannot parse json credentials")
+			}
+		}
+	} else {
+		return nil, errors.New("failed to get azure-creds credentials")
+	}
+
+	return azureCreds, nil
+}
+
+// AzureQuery is a concrete implementation of the AzureQueryInterface
+// that interacts with Azure Resource Graph API.
+type AzureQuery struct{}
+
+func (a *AzureQuery) azQuery(ctx context.Context, azureCreds map[string]string, in *v1beta1.Input) (armresourcegraph.ClientResourcesResponse, error) {
+	tenantID := azureCreds["tenantId"]
+	clientID := azureCreds["clientId"]
+	clientSecret := azureCreds["clientSecret"]
+	subscriptionID := azureCreds["subscriptionId"]
+
+	// To configure DefaultAzureCredential to authenticate a user-assigned managed identity,
+	// set the environment variable AZURE_CLIENT_ID to the identity's client ID.
+
+	cred, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, nil)
+	if err != nil {
+		return armresourcegraph.ClientResourcesResponse{}, errors.Wrap(err, "failed to obtain credentials")
+	}
+
+	// Create and authorize a ResourceGraph client
+	client, err := armresourcegraph.NewClient(cred, nil)
+	if err != nil {
+		return armresourcegraph.ClientResourcesResponse{}, errors.Wrap(err, "failed to create client")
+	}
+
+	queryRequest := armresourcegraph.QueryRequest{
+		Query: to.Ptr(in.Query),
+	}
+
+	if len(subscriptionID) > 0 {
+		queryRequest.Subscriptions = []*string{to.Ptr(subscriptionID)}
+	}
+
+	if len(in.ManagementGroups) > 0 {
+		queryRequest.ManagementGroups = in.ManagementGroups
+	}
+
+	// Create the query request, Run the query and get the results.
+	results, err := client.Resources(ctx, queryRequest, nil)
+	if err != nil {
+		return armresourcegraph.ClientResourcesResponse{}, errors.Wrap(err, "failed to finish the request")
+	}
+	return results, nil
+}

fn_test.go

@@ -4,8 +4,11 @@ import (
 	"context"
 	"testing"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/upboundcare/function-azresourcegraph/input/v1beta1"
 	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/durationpb"
 
@@ -15,8 +18,30 @@ import (
 	"github.com/crossplane/function-sdk-go/response"
 )
 
+type MockAzureQuery struct {
+	AzQueryFunc func(ctx context.Context, azureCreds map[string]string, in *v1beta1.Input) (armresourcegraph.ClientResourcesResponse, error)
+}
+
+func (m *MockAzureQuery) azQuery(ctx context.Context, azureCreds map[string]string, in *v1beta1.Input) (armresourcegraph.ClientResourcesResponse, error) {
+	return m.AzQueryFunc(ctx, azureCreds, in)
+}
+
 func TestRunFunction(t *testing.T) {
 
+	var (
+		xr    = `{"apiVersion":"example.org/v1","kind":"XR","metadata":{"name":"cool-xr"},"spec":{"count":2}}`
+		creds = &fnv1.CredentialData{
+			Data: map[string][]byte{
+				"credentials": []byte(`{
+"clientId": "test-cliend-id",
+"clientSecret": "test-client-secret",
+"subscriptionId": "test-subscription-id",
+"tenantId": "test-tenant-id"
+}`),
+			},
+		}
+	)
+
 	type args struct {
 		ctx context.Context
 		req *fnv1.RunFunctionRequest
@@ -82,11 +107,88 @@ func TestRunFunction(t *testing.T) {
 				},
 			},
 		},
+		"ShouldUpdateXRStatus": {
+			reason: "The Function should update XR status",
+			args: args{
+				ctx: context.Background(),
+				req: &fnv1.RunFunctionRequest{
+					Meta: &fnv1.RequestMeta{Tag: "hello"},
+					Input: resource.MustStructJSON(`{
+						"apiVersion": "azresourcegraph.fn.crossplane.io/v1alpha1",
+						"kind": "Input",
+						"query": "Resources| count",
+						"managementGroups": ["test"]
+					}`),
+					Observed: &fnv1.State{
+						Composite: &fnv1.Resource{
+							Resource: resource.MustStructJSON(xr),
+						},
+					},
+					Credentials: map[string]*fnv1.Credentials{
+						"azure-creds": {
+							Source: &fnv1.Credentials_CredentialData{CredentialData: creds},
+						},
+					},
+				},
+			},
+			want: want{
+				rsp: &fnv1.RunFunctionResponse{
+					Meta: &fnv1.ResponseMeta{Tag: "hello", Ttl: durationpb.New(response.DefaultTTL)},
+					Conditions: []*fnv1.Condition{
+						{
+							Type:   "FunctionSuccess",
+							Status: fnv1.Status_STATUS_CONDITION_TRUE,
+							Reason: "Success",
+							Target: fnv1.Target_TARGET_COMPOSITE_AND_CLAIM.Enum(),
+						},
+					},
+					Results: []*fnv1.Result{
+						{
+							Severity: fnv1.Severity_SEVERITY_NORMAL,
+							Message:  `Query: "Resources| count"`,
+							Target:   fnv1.Target_TARGET_COMPOSITE.Enum(),
+						},
+					},
+					Desired: &fnv1.State{
+						Composite: &fnv1.Resource{
+							Resource: resource.MustStructJSON(`{
+"apiVersion": "example.org/v1",
+"kind": "XR",
+"status": {
+	"azResourceGraphQueryResult": [
+		{
+			"resource": "mock-resource"
+		}
+	]
+}}`),
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
-			f := &Function{log: logging.NewNopLogger()}
+			// Mocking the azQuery function to return a successful result
+			mockQuery := &MockAzureQuery{
+				AzQueryFunc: func(_ context.Context, _ map[string]string, _ *v1beta1.Input) (armresourcegraph.ClientResourcesResponse, error) {
+					return armresourcegraph.ClientResourcesResponse{
+						QueryResponse: armresourcegraph.QueryResponse{
+							Count:           to.Ptr(int64(1)),
+							Data:            []map[string]interface{}{{"resource": "mock-resource"}}, // Mock data
+							ResultTruncated: to.Ptr(armresourcegraph.ResultTruncatedFalse),
+							TotalRecords:    to.Ptr(int64(1)),
+							Facets:          nil,
+							SkipToken:       nil,
+						},
+					}, nil
+				},
+			}
+			f := &Function{
+				azureQuery: mockQuery,
+				log:        logging.NewNopLogger(),
+			}
 			rsp, err := f.RunFunction(tc.args.ctx, tc.args.req)
 
 			if diff := cmp.Diff(tc.want.rsp, rsp, protocmp.Transform()); diff != "" {