new-api/service/asset_chinamobile_test.go

package service

import (
	"context"
	"errors"
	"net/http"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/model"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	cmerrs "gitlab.ecloud.com/ecloud/ecloudsdkcore/errs"
	cmmodel "gitlab.ecloud.com/ecloud/ecloudsdkmaas/model"
)

func TestChinaMobileAssetAdapterCreateAssetMapsRequestAndResponse(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "")

	fake := &fakeChinaMobileAssetSDKClient{}
	adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient {
		assert.Equal(t, "ak", credential.AK)
		assert.Equal(t, "sk", credential.SK)
		assert.Equal(t, defaultChinaMobileAssetPoolID, credential.PoolID)
		return fake
	}}

	spec, ok := ParseAssetAction("CreateAsset")
	require.True(t, ok)
	resp, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "video-generation-key"}, AssetRequest{
		Action:  spec,
		Version: "2024-01-01",
		Body: map[string]any{
			"GroupId":   "group-1",
			"Name":      "asset",
			"URL":       "https://example.com/a.png",
			"AssetType": "Image",
		},
	})

	require.Nil(t, assetErr)
	require.NotNil(t, resp)
	require.NotNil(t, fake.createAssetRequest)
	assert.Equal(t, "group-1", *fake.createAssetRequest.CreateAssetBody.GroupId)
	assert.Equal(t, "asset", *fake.createAssetRequest.CreateAssetBody.AssetName)
	assert.Equal(t, "https://example.com/a.png", *fake.createAssetRequest.CreateAssetBody.AssetUrl)
	assert.Equal(t, cmmodel.CreateAssetBodyAssetTypeEnumImage, *fake.createAssetRequest.CreateAssetBody.AssetType)
	assert.Contains(t, string(resp.Body), `"RequestId":"req-1"`)
	assert.Contains(t, string(resp.Body), `"Result":"asset-1"`)
}

func TestChinaMobileAssetCredentialFromEnvDefaultsCenterPool(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "")

	credential, err := chinaMobileAssetCredentialFromEnv()

	require.NoError(t, err)
	assert.Equal(t, "ak", credential.AK)
	assert.Equal(t, "sk", credential.SK)
	assert.Equal(t, defaultChinaMobileAssetPoolID, credential.PoolID)
}

func TestChinaMobileAssetCredentialFromEnvSupportsPoolID(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "CIDC-RP-29")

	credential, err := chinaMobileAssetCredentialFromEnv()

	require.NoError(t, err)
	assert.Equal(t, "ak", credential.AK)
	assert.Equal(t, "sk", credential.SK)
	assert.Equal(t, "CIDC-RP-29", credential.PoolID)
}

func TestChinaMobileAssetCredentialFromEnvRequiresAKAndSK(t *testing.T) {
	t.Setenv(chinaMobileAssetAKEnv, "")
	t.Setenv(chinaMobileAssetSKEnv, "")
	t.Setenv(chinaMobileAssetPoolIDEnv, "")

	_, err := chinaMobileAssetCredentialFromEnv()

	require.Error(t, err)
	assert.Contains(t, err.Error(), chinaMobileAssetAKEnv)
	assert.Contains(t, err.Error(), chinaMobileAssetSKEnv)
}

func TestChinaMobileAssetAdapterAllOfficialActions(t *testing.T) {
	cases := []struct {
		action string
		body   map[string]any
		assert func(t *testing.T, fake *fakeChinaMobileAssetSDKClient)
	}{
		{"CreateAssetGroup", map[string]any{"Name": "g", "GroupType": "AIGC"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) {
			require.NotNil(t, fake.createAssetGroupRequest)
		}},
		{"CreateAsset", map[string]any{"GroupId": "g", "Name": "n", "URL": "https://example.com/a.png", "AssetType": "Image"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) { require.NotNil(t, fake.createAssetRequest) }},
		{"ListAssetGroups", map[string]any{"Filter": map[string]any{"GroupType": "AIGC"}, "PageNumber": float64(1), "PageSize": float64(10)}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) {
			require.NotNil(t, fake.listAssetGroupsRequest)
		}},
		{"ListAssets", map[string]any{"Filter": map[string]any{"GroupType": "AIGC"}, "PageNumber": float64(1), "PageSize": float64(10)}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) { require.NotNil(t, fake.listAssetsRequest) }},
		{"GetAsset", map[string]any{"Id": "asset-1"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) {
			require.NotNil(t, fake.getAssetRequest)
			assert.Equal(t, "asset-1", *fake.getAssetRequest.GetAssetPath.AssetId)
		}},
		{"GetAssetGroup", map[string]any{"Id": "group-1"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) {
			require.NotNil(t, fake.getAssetGroupRequest)
			assert.Equal(t, "group-1", *fake.getAssetGroupRequest.GetAssetGroupPath.GroupId)
		}},
		{"UpdateAsset", map[string]any{"Id": "asset-1", "Name": "n"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) { require.NotNil(t, fake.updateAssetRequest) }},
		{"UpdateAssetGroup", map[string]any{"Id": "group-1", "Name": "g"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) {
			require.NotNil(t, fake.updateAssetGroupRequest)
		}},
		{"DeleteAsset", map[string]any{"Id": "asset-1"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) { require.NotNil(t, fake.deleteAssetRequest) }},
		{"DeleteAssetGroup", map[string]any{"Id": "group-1"}, func(t *testing.T, fake *fakeChinaMobileAssetSDKClient) {
			require.NotNil(t, fake.deleteAssetGroupRequest)
		}},
	}

	for _, tc := range cases {
		t.Run(tc.action, func(t *testing.T) {
			setChinaMobileAssetTestEnv(t, "")

			fake := &fakeChinaMobileAssetSDKClient{}
			adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient {
				return fake
			}}
			spec, ok := ParseAssetAction(tc.action)
			require.True(t, ok)

			resp, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "video-generation-key"}, AssetRequest{
				Action:  spec,
				Version: "2024-01-01",
				Body:    tc.body,
			})

			require.Nil(t, assetErr)
			require.NotNil(t, resp)
			tc.assert(t, fake)
			assert.Contains(t, string(resp.Body), `"ResponseMetadata"`)
		})
	}
}

func TestChinaMobileAssetAdapterRejectsUnsupportedCompatibilityOptions(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "")

	adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient {
		return &fakeChinaMobileAssetSDKClient{}
	}}
	spec, ok := ParseAssetAction("ListAssets")
	require.True(t, ok)

	cases := []map[string]any{
		{"ProjectName": "prod"},
		{"SortBy": "Name"},
		{"SortOrder": "Asc"},
	}
	for _, body := range cases {
		_, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "cm-key"}, AssetRequest{
			Action: spec,
			Body:   body,
		})

		require.NotNil(t, assetErr)
		assert.Equal(t, AssetErrorOperationNotSupported, assetErr.Type)
	}
}

func TestChinaMobileAssetAdapterMissingIdReturnsInvalidRequest(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "")

	adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient {
		return &fakeChinaMobileAssetSDKClient{}
	}}
	spec, ok := ParseAssetAction("GetAsset")
	require.True(t, ok)

	_, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "video-generation-key"}, AssetRequest{
		Action: spec,
		Body:   map[string]any{},
	})

	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorInvalidRequest, assetErr.Type)
	assert.Contains(t, assetErr.Message, "Id is required")
}

func TestChinaMobileAssetAdapterStateErrorMapsToUpstreamError(t *testing.T) {
	_, assetErr := normalizeChinaMobileAssetResponse(AssetActionSpec{Action: "CreateAsset", Operation: AssetOperationAssetCreate}, "2024-01-01", []byte(`{"requestId":"req","state":"ERROR","errorCode":"Bad","errorMessage":"bad request"}`))

	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorUpstream, assetErr.Type)
	assert.Equal(t, "bad request", assetErr.Message)
}

func TestChinaMobileAssetAdapterRequiresGroupTypeForListAssets(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "")

	adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient {
		return &fakeChinaMobileAssetSDKClient{}
	}}
	spec, ok := ParseAssetAction("ListAssets")
	require.True(t, ok)

	_, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "video-generation-key"}, AssetRequest{
		Action: spec,
		Body:   map[string]any{"PageNumber": 1, "PageSize": 10},
	})

	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorInvalidRequest, assetErr.Type)
	assert.Equal(t, http.StatusBadRequest, assetErr.HTTPStatus)
}

func TestChinaMobileAssetDeleteFalseMapsToUpstreamError(t *testing.T) {
	_, assetErr := normalizeChinaMobileAssetResponse(AssetActionSpec{Action: "DeleteAsset", Operation: AssetOperationAssetDelete, Delete: true}, "2024-01-01", []byte(`{"requestId":"req","state":"OK","body":false}`))

	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorUpstream, assetErr.Type)
}

func TestChinaMobileAssetUnknownStateMapsToUpstreamError(t *testing.T) {
	_, assetErr := normalizeChinaMobileAssetResponse(AssetActionSpec{Action: "CreateAsset", Operation: AssetOperationAssetCreate}, "2024-01-01", []byte(`{"requestId":"req","state":"","body":"asset-1"}`))

	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorUpstream, assetErr.Type)
}

func TestChinaMobileAssetSDKBadRequestMapsToInvalidRequest(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "")

	adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient {
		return &errorChinaMobileAssetSDKClient{
			fakeChinaMobileAssetSDKClient: fakeChinaMobileAssetSDKClient{},
			err:                           cmerrs.NewServerResponseError("bad request", nil, http.StatusBadRequest, nil, `{"errorMessage":"invalid asset"}`),
		}
	}}
	spec, ok := ParseAssetAction("CreateAsset")
	require.True(t, ok)

	_, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "video-generation-key"}, AssetRequest{Action: spec, Body: validChinaMobileCreateAssetBody()})

	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorInvalidRequest, assetErr.Type)
	assert.Equal(t, http.StatusBadRequest, assetErr.HTTPStatus)
}

func TestChinaMobileAssetSDKServerErrorRemainsUpstreamError(t *testing.T) {
	setChinaMobileAssetTestEnv(t, "")

	adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient {
		return &errorChinaMobileAssetSDKClient{
			fakeChinaMobileAssetSDKClient: fakeChinaMobileAssetSDKClient{},
			err:                           cmerrs.NewServerResponseError("server error", nil, http.StatusInternalServerError, nil, ""),
		}
	}}
	spec, ok := ParseAssetAction("CreateAsset")
	require.True(t, ok)

	_, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "video-generation-key"}, AssetRequest{Action: spec, Body: validChinaMobileCreateAssetBody()})

	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorUpstream, assetErr.Type)
	assert.Equal(t, http.StatusBadGateway, assetErr.HTTPStatus)
}

func TestCallChinaMobileAssetSDKReturnsWhenContextCancelled(t *testing.T) {
	client := &blockingChinaMobileAssetSDKClient{fakeChinaMobileAssetSDKClient: fakeChinaMobileAssetSDKClient{}, release: make(chan struct{})}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	spec, ok := ParseAssetAction("CreateAsset")
	require.True(t, ok)

	started := time.Now()
	_, err := callChinaMobileAssetSDK(ctx, client, AssetRequest{Action: spec, Body: validChinaMobileCreateAssetBody()})

	require.ErrorIs(t, err, context.Canceled)
	assert.Less(t, time.Since(started), 500*time.Millisecond)
	close(client.release)
}

func setChinaMobileAssetTestEnv(t *testing.T, poolID string) {
	t.Helper()
	t.Setenv(chinaMobileAssetAKEnv, "ak")
	t.Setenv(chinaMobileAssetSKEnv, "sk")
	t.Setenv(chinaMobileAssetPoolIDEnv, poolID)
}

type blockingChinaMobileAssetSDKClient struct {
	fakeChinaMobileAssetSDKClient
	release chan struct{}
}

type errorChinaMobileAssetSDKClient struct {
	fakeChinaMobileAssetSDKClient
	err error
}

func (f *errorChinaMobileAssetSDKClient) CreateAsset(request *cmmodel.CreateAssetRequest) (*cmmodel.CreateAssetResponse, error) {
	if f.err == nil {
		return nil, errors.New("missing test error")
	}
	return nil, f.err
}

func (f *blockingChinaMobileAssetSDKClient) CreateAsset(request *cmmodel.CreateAssetRequest) (*cmmodel.CreateAssetResponse, error) {
	<-f.release
	return f.fakeChinaMobileAssetSDKClient.CreateAsset(request)
}

func validChinaMobileCreateAssetBody() map[string]any {
	return map[string]any{
		"GroupId":   "group-1",
		"Name":      "asset",
		"URL":       "https://example.com/asset.png",
		"AssetType": "Image",
	}
}

type fakeChinaMobileAssetSDKClient struct {
	createAssetRequest      *cmmodel.CreateAssetRequest
	listAssetsRequest       *cmmodel.ListAssetsRequest
	getAssetRequest         *cmmodel.GetAssetRequest
	updateAssetRequest      *cmmodel.UpdateAssetRequest
	deleteAssetRequest      *cmmodel.DeleteAssetRequest
	createAssetGroupRequest *cmmodel.CreateAssetGroupRequest
	listAssetGroupsRequest  *cmmodel.ListAssetGroupsRequest
	getAssetGroupRequest    *cmmodel.GetAssetGroupRequest
	updateAssetGroupRequest *cmmodel.UpdateAssetGroupRequest
	deleteAssetGroupRequest *cmmodel.DeleteAssetGroupRequest
}

func (f *fakeChinaMobileAssetSDKClient) CreateAsset(request *cmmodel.CreateAssetRequest) (*cmmodel.CreateAssetResponse, error) {
	f.createAssetRequest = request
	return (&cmmodel.CreateAssetResponse{}).SetRequestId("req-1").SetState("OK").SetBody("asset-1"), nil
}

func (f *fakeChinaMobileAssetSDKClient) ListAssets(request *cmmodel.ListAssetsRequest) (*cmmodel.ListAssetsResponse, error) {
	f.listAssetsRequest = request
	item := cmmodel.ListAssetsResponseData{}
	item.SetAssetId("asset-1").SetGroupId("group-1").SetAssetName("n").SetStatus(cmmodel.ListAssetsResponseDataStatusEnumActive)
	body := (&cmmodel.ListAssetsResponseBody{}).SetData([]cmmodel.ListAssetsResponseData{item}).SetTotal(1)
	return (&cmmodel.ListAssetsResponse{}).SetRequestId("req").SetState("OK").SetBody(body), nil
}

func (f *fakeChinaMobileAssetSDKClient) GetAsset(request *cmmodel.GetAssetRequest) (*cmmodel.GetAssetResponse, error) {
	f.getAssetRequest = request
	body := (&cmmodel.GetAssetResponseBody{}).SetAssetId("asset-1").SetGroupId("group-1").SetAssetName("n").SetStatus(cmmodel.GetAssetResponseBodyStatusEnumActive)
	return (&cmmodel.GetAssetResponse{}).SetRequestId("req").SetState("OK").SetBody(body), nil
}

func (f *fakeChinaMobileAssetSDKClient) UpdateAsset(request *cmmodel.UpdateAssetRequest) (*cmmodel.UpdateAssetResponse, error) {
	f.updateAssetRequest = request
	body := (&cmmodel.UpdateAssetResponseBody{}).SetAssetId("asset-1").SetGroupId("group-1").SetAssetName("n").SetStatus(cmmodel.UpdateAssetResponseBodyStatusEnumActive)
	return (&cmmodel.UpdateAssetResponse{}).SetRequestId("req").SetState("OK").SetBody(body), nil
}

func (f *fakeChinaMobileAssetSDKClient) DeleteAsset(request *cmmodel.DeleteAssetRequest) (*cmmodel.DeleteAssetResponse, error) {
	f.deleteAssetRequest = request
	return (&cmmodel.DeleteAssetResponse{}).SetRequestId("req").SetState("OK").SetBody(true), nil
}

func (f *fakeChinaMobileAssetSDKClient) CreateAssetGroup(request *cmmodel.CreateAssetGroupRequest) (*cmmodel.CreateAssetGroupResponse, error) {
	f.createAssetGroupRequest = request
	body := (&cmmodel.CreateAssetGroupResponseBody{}).SetGroupId("group-1").SetGroupName("g").SetGroupType(cmmodel.CreateAssetGroupResponseBodyGroupTypeEnumAigc)
	return (&cmmodel.CreateAssetGroupResponse{}).SetRequestId("req").SetState("OK").SetBody(body), nil
}

func (f *fakeChinaMobileAssetSDKClient) ListAssetGroups(request *cmmodel.ListAssetGroupsRequest) (*cmmodel.ListAssetGroupsResponse, error) {
	f.listAssetGroupsRequest = request
	item := cmmodel.ListAssetGroupsResponseData{}
	item.SetGroupId("group-1").SetGroupName("g").SetGroupType(cmmodel.ListAssetGroupsResponseDataGroupTypeEnumAigc)
	body := (&cmmodel.ListAssetGroupsResponseBody{}).SetData([]cmmodel.ListAssetGroupsResponseData{item}).SetTotal(1)
	return (&cmmodel.ListAssetGroupsResponse{}).SetRequestId("req").SetState("OK").SetBody(body), nil
}

func (f *fakeChinaMobileAssetSDKClient) GetAssetGroup(request *cmmodel.GetAssetGroupRequest) (*cmmodel.GetAssetGroupResponse, error) {
	f.getAssetGroupRequest = request
	body := (&cmmodel.GetAssetGroupResponseBody{}).SetGroupId("group-1").SetGroupName("g").SetGroupType(cmmodel.GetAssetGroupResponseBodyGroupTypeEnumAigc)
	return (&cmmodel.GetAssetGroupResponse{}).SetRequestId("req").SetState("OK").SetBody(body), nil
}

func (f *fakeChinaMobileAssetSDKClient) UpdateAssetGroup(request *cmmodel.UpdateAssetGroupRequest) (*cmmodel.UpdateAssetGroupResponse, error) {
	f.updateAssetGroupRequest = request
	body := (&cmmodel.UpdateAssetGroupResponseBody{}).SetGroupId("group-1").SetGroupName("g").SetGroupType(cmmodel.UpdateAssetGroupResponseBodyGroupTypeEnumAigc)
	return (&cmmodel.UpdateAssetGroupResponse{}).SetRequestId("req").SetState("OK").SetBody(body), nil
}

func (f *fakeChinaMobileAssetSDKClient) DeleteAssetGroup(request *cmmodel.DeleteAssetGroupRequest) (*cmmodel.DeleteAssetGroupResponse, error) {
	f.deleteAssetGroupRequest = request
	return (&cmmodel.DeleteAssetGroupResponse{}).SetRequestId("req").SetState("OK").SetBody(true), nil
}