go/chat/attachments/uploader_test.go - master - githubmirror/client

go/chat/attachments/uploader_test.go
370 строк10 KB
zoom-ua
remove golang.org/x/net/context in favor of context (#28708)
17 дек 2025, 18:24
17 дек 2025, 18:2467b0ffe
package attachments

import (
	"context"
	"errors"
	"io"
	"path/filepath"
	"runtime"
	"testing"
	"time"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/types"
	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/stretchr/testify/require"

	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
)

type mockStore struct {
	Store
	uploadFn func(context.Context, *UploadTask) (chat1.Asset, error)
}

func (m *mockStore) UploadAsset(ctx context.Context, task *UploadTask, encryptedOut io.Writer) (chat1.Asset, error) {
	return m.uploadFn(ctx, task)
}

type mockRemote struct {
	chat1.RemoteInterface
}

func (r mockRemote) GetS3Params(context.Context, chat1.GetS3ParamsArg) (chat1.S3Params, error) {
	return chat1.S3Params{}, nil
}

func (r mockRemote) S3Sign(context.Context, chat1.S3SignArg) ([]byte, error) {
	return nil, nil
}

type mockActivityNotifier struct {
	types.ActivityNotifier
	startCh chan chat1.OutboxID
}

func newMockActivityNotifier() *mockActivityNotifier {
	return &mockActivityNotifier{
		startCh: make(chan chat1.OutboxID, 1000),
	}
}

func (a *mockActivityNotifier) AttachmentUploadStart(ctx context.Context, uid gregor1.UID,
	convID chat1.ConversationID, outboxID chat1.OutboxID,
) {
	a.startCh <- outboxID
}

func (a *mockActivityNotifier) AttachmentUploadProgress(ctx context.Context, uid gregor1.UID,
	convID chat1.ConversationID, outboxID chat1.OutboxID, bytesComplete, bytesTotal int64) {
}

type mockDeliverer struct {
	types.MessageDeliverer
	forceCh chan struct{}
}

func newMockDeliverer() *mockDeliverer {
	return &mockDeliverer{
		forceCh: make(chan struct{}, 1000),
	}
}

func (m *mockDeliverer) ForceDeliverLoop(context.Context) {
	m.forceCh <- struct{}{}
}

func (m *mockDeliverer) Stop(context.Context) chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
}

type mockInboxSource struct {
	types.InboxSource
}

func (m mockInboxSource) ReadUnverified(ctx context.Context, uid gregor1.UID,
	dataSource types.InboxSourceDataSourceTyp, rquery *chat1.GetInboxQuery,
) (types.Inbox, error) {
	return types.Inbox{
		ConvsUnverified: []types.RemoteConversation{
			{
				Conv: chat1.Conversation{
					Metadata: chat1.ConversationMetadata{
						ConversationID: chat1.ConversationID([]byte{0, 1, 0}),
						IdTriple: chat1.ConversationIDTriple{
							TopicType: chat1.TopicType_CHAT,
						},
					},
				},
			},
		},
	}, nil
}

func (m mockInboxSource) Stop(context.Context) chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
}

func TestAttachmentUploader(t *testing.T) {
	world := kbtest.NewChatMockWorld(t, "uploader", 1)
	defer world.Cleanup()

	u := world.GetUsers()[0]
	uid := gregor1.UID(u.User.GetUID().ToBytes())
	tc := world.Tcs[u.Username]
	g := globals.NewContext(tc.G, tc.ChatG)
	notifier := newMockActivityNotifier()
	store := &mockStore{}
	ri := mockRemote{}
	deliverer := newMockDeliverer()
	g.AttachmentURLSrv = types.DummyAttachmentHTTPSrv{}
	g.InboxSource = mockInboxSource{}
	g.ActivityNotifier = notifier
	g.MessageDeliverer = deliverer
	getRi := func() chat1.RemoteInterface { return ri }
	cacheSize := 1
	uploader := NewUploader(g, store, NewS3Signer(getRi), getRi, cacheSize)
	convID := chat1.ConversationID([]byte{0, 1, 0})
	md, err := libkb.RandBytes(10)
	require.NoError(t, err)

	uploadStartCheck := func(shouldHappen bool, outboxID chat1.OutboxID) {
		if shouldHappen {
			select {
			case obid := <-notifier.startCh:
				require.Equal(t, outboxID, obid)
			case <-time.After(20 * time.Second):
				require.Fail(t, "no start")
			}
		} else {
			select {
			case <-notifier.startCh:
				require.Fail(t, "start not supposed to happen")
			default:
			}
		}
	}
	deliverCheck := func(shouldHappen bool) {
		if shouldHappen {
			select {
			case <-deliverer.forceCh:
			case <-time.After(20 * time.Second):
				require.Fail(t, "no start")
			}
		} else {
			select {
			case <-deliverer.forceCh:
				require.Fail(t, "start not supposed to happen")
			default:
			}
		}
	}
	successCheck := func(cb types.AttachmentUploaderResultCb) {
		ch := cb.Wait()
		select {
		case res := <-ch:
			require.Nil(t, res.Error)
			require.Equal(t, md, res.Metadata)
			require.NotNil(t, res.Preview)
			require.Equal(t, "image/jpeg", res.Preview.MimeType)
			require.Equal(t, "image/jpeg", res.Object.MimeType)
		case <-time.After(20 * time.Second):
			require.Fail(t, "no upload")
		}
	}
	// On non darwin we don't covert the heic.
	successCheckNoHeicConvert := func(cb types.AttachmentUploaderResultCb) {
		ch := cb.Wait()
		select {
		case res := <-ch:
			require.Nil(t, res.Error)
			require.Equal(t, md, res.Metadata)
			require.Nil(t, res.Preview)
			require.Equal(t, "image/heif", res.Object.MimeType)
		case <-time.After(20 * time.Second):
			require.Fail(t, "no upload")
		}
	}

	successCheckEmpty := func(cb types.AttachmentUploaderResultCb) {
		ch := cb.Wait()
		select {
		case res := <-ch:
			require.Nil(t, res.Error)
			require.Equal(t, md, res.Metadata)
			require.Nil(t, res.Preview)
			require.Equal(t, "", res.Object.MimeType)
		case <-time.After(20 * time.Second):
			require.Fail(t, "no upload")
		}
	}

	// Basic test to see if it works
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, nil
	}

	outboxID, err := storage.NewOutboxID()
	require.NoError(t, err)
	filename := "../testdata/empty.txt"
	resChan, err := uploader.Register(context.TODO(), uid, convID, outboxID, "empty", filename, md, nil)
	require.NoError(t, err)
	deliverCheck(true)
	uploadStartCheck(true, outboxID)
	successCheckEmpty(resChan)

	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	filename = "../testdata/mysql.heic"
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "mysql", filename, md, nil)
	require.NoError(t, err)
	deliverCheck(true)
	uploadStartCheck(true, outboxID)
	if runtime.GOOS == "darwin" {
		successCheck(resChan)
	} else {
		successCheckNoHeicConvert(resChan)
	}

	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	filename = "../testdata/ship.jpg"
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	deliverCheck(true)
	uploadStartCheck(true, outboxID)
	successCheck(resChan)

	// Broken store
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, errors.New("i dont work")
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	select {
	case res := <-resChan.Wait():
		require.NotNil(t, res.Error)
	case <-time.After(20 * time.Second):
		require.Fail(t, "no upload")
	}
	deliverCheck(true)

	// block until the upload is marked as done
	for count := 0; count <= 5; count++ {
		uploader.Lock()
		upload, ok := uploader.uploads[outboxID.String()]
		uploader.Unlock()
		if !ok && upload == nil {
			break
		}
		time.Sleep(time.Millisecond * 200)
		if count == 5 {
			require.Fail(t, "upload not marked as done")
		}
		t.Logf("upload not done, checking again")
	}
	t.Logf("upload done")

	// Retry after fixing store
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Retry(context.TODO(), outboxID)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	successCheck(resChan)
	deliverCheck(true)

	// Slow store to test concurrent retry
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	slowCh := make(chan struct{})
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		<-slowCh
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	deliverCheck(false)
	select {
	case <-resChan.Wait():
		require.Fail(t, "no res")
	default:
	}
	retryChan, err := uploader.Retry(context.TODO(), outboxID)
	require.NoError(t, err)
	uploadStartCheck(false, outboxID)
	close(slowCh)
	deliverCheck(true)
	// Should get results on both of these
	successCheck(retryChan)
	successCheck(resChan)

	uploader.Complete(context.TODO(), outboxID)
	_, _, err = uploader.Status(context.TODO(), outboxID)
	require.Error(t, err)

	// Test cancel
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	slowCh = make(chan struct{})
	store.uploadFn = func(ctx context.Context, task *UploadTask) (chat1.Asset, error) {
		select {
		case <-slowCh:
		case <-ctx.Done():
			return chat1.Asset{}, ctx.Err()
		}
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	deliverCheck(false)
	select {
	case <-resChan.Wait():
		require.Fail(t, "no res")
	default:
	}
	require.NoError(t, uploader.Cancel(context.TODO(), outboxID))
	_, _, err = uploader.Status(context.TODO(), outboxID)
	require.Error(t, err)
	res := <-resChan.Wait()
	require.NotNil(t, res.Error)

	// verify uploadedPreviewsDir respects the cache size
	baseDir := uploader.getBaseDir()
	uploadedPreviews, err := filepath.Glob(filepath.Join(baseDir, uploadedPreviewsDir, "*"))
	require.NoError(t, err)
	require.Len(t, uploadedPreviews, 1)

	// verify uploadedFullsDir is respects the cache size
	uploadedFulls, err := filepath.Glob(filepath.Join(baseDir, uploadedFullsDir, "*"))
	require.NoError(t, err)
	require.Len(t, uploadedFulls, 1)
	mctx := kbtest.NewMetaContextForTest(*tc)

	// verify db nuke
	_, err = g.LocalDb.Nuke()
	require.NoError(t, err)
	err = uploader.OnDbNuke(mctx)
	require.NoError(t, err)

	uploadedPreviews, err = filepath.Glob(filepath.Join(baseDir, uploadedPreviewsDir, "*"))
	require.NoError(t, err)
	require.Zero(t, len(uploadedPreviews))

	uploadedFulls, err = filepath.Glob(filepath.Join(baseDir, uploadedFullsDir, "*"))
	require.NoError(t, err)
	require.Zero(t, len(uploadedFulls))
}