image_task_test.go 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. package service
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "net/http"
  7. "testing"
  8. "time"
  9. "github.com/stretchr/testify/require"
  10. )
  11. type imageTaskMemoryStore struct {
  12. task *ImageTaskRecord
  13. ttl time.Duration
  14. saveErr error
  15. getErr error
  16. }
  17. func (s *imageTaskMemoryStore) Save(_ context.Context, task *ImageTaskRecord, ttl time.Duration) error {
  18. if s.saveErr != nil {
  19. return s.saveErr
  20. }
  21. copy := *task
  22. s.task = &copy
  23. s.ttl = ttl
  24. return nil
  25. }
  26. func (s *imageTaskMemoryStore) Get(_ context.Context, _ string) (*ImageTaskRecord, error) {
  27. if s.getErr != nil {
  28. return nil, s.getErr
  29. }
  30. if s.task == nil {
  31. return nil, ErrImageTaskNotFound
  32. }
  33. copy := *s.task
  34. return &copy, nil
  35. }
  36. func TestImageTaskServiceLifecycleAndOwnership(t *testing.T) {
  37. store := &imageTaskMemoryStore{}
  38. svc := NewImageTaskServiceWithOptions(store, time.Hour, 10*time.Minute)
  39. owner := ImageTaskOwner{UserID: 7, APIKeyID: 9}
  40. created, err := svc.Create(context.Background(), owner)
  41. require.NoError(t, err)
  42. require.Equal(t, ImageTaskStatusProcessing, created.Status)
  43. require.Equal(t, created.ID, created.TaskID)
  44. require.Equal(t, "image.generation.task", created.Object)
  45. require.Equal(t, time.Hour, store.ttl)
  46. require.Equal(t, owner.UserID, store.task.UserID)
  47. require.Equal(t, owner.APIKeyID, store.task.APIKeyID)
  48. _, err = svc.Get(context.Background(), ImageTaskOwner{UserID: 7, APIKeyID: 10}, created.ID)
  49. require.ErrorIs(t, err, ErrImageTaskNotFound)
  50. result := json.RawMessage(`{"created":123,"data":[{"url":"https://example.test/image.png"}]}`)
  51. require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, result))
  52. completed, err := svc.Get(context.Background(), owner, created.ID)
  53. require.NoError(t, err)
  54. require.Equal(t, ImageTaskStatusCompleted, completed.Status)
  55. require.Equal(t, http.StatusOK, completed.HTTPStatus)
  56. require.Equal(t, "https://example.test/image.png", completed.ImageURL)
  57. require.JSONEq(t, string(result), string(completed.Result))
  58. require.NotNil(t, completed.CompletedAt)
  59. }
  60. func TestImageTaskServiceInvalidResultBecomesFailed(t *testing.T) {
  61. store := &imageTaskMemoryStore{}
  62. svc := NewImageTaskServiceWithOptions(store, time.Hour, time.Minute)
  63. created, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2})
  64. require.NoError(t, err)
  65. require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, json.RawMessage(`not-json`)))
  66. got, err := svc.Get(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2}, created.ID)
  67. require.NoError(t, err)
  68. require.Equal(t, ImageTaskStatusFailed, got.Status)
  69. require.Equal(t, http.StatusBadGateway, got.HTTPStatus)
  70. require.Contains(t, string(got.Error), "non-JSON")
  71. }
  72. func TestImageTaskServiceMapsStoreFailures(t *testing.T) {
  73. store := &imageTaskMemoryStore{saveErr: errors.New("redis down")}
  74. svc := NewImageTaskService(store)
  75. _, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2})
  76. require.ErrorIs(t, err, ErrImageTaskUnavailable)
  77. }