feat: supports tiktok apify capabilities

Makefile

@@ -77,7 +77,7 @@ test-twitter: docker-build-test
 	@docker run --user root $(ENV_FILE_ARG) -v $(PWD)/.masa:/home/masa -v $(PWD)/coverage:/app/coverage --rm --workdir /app -e DATA_DIR=/home/masa $(TEST_IMAGE) go test -v ./internal/jobs/twitter_test.go ./internal/jobs/jobs_suite_test.go
 
 test-tiktok: docker-build-test
-	@docker run --user root $(ENV_FILE_ARG) -v $(PWD)/.masa:/home/masa -v $(PWD)/coverage:/app/coverage --rm --workdir /app -e DATA_DIR=/home/masa $(TEST_IMAGE) go test -v ./internal/jobs/tiktok_transcription_test.go ./internal/jobs/jobs_suite_test.go
+	@docker run --user root $(ENV_FILE_ARG) -v $(PWD)/.masa:/home/masa -v $(PWD)/coverage:/app/coverage --rm --workdir /app -e DATA_DIR=/home/masa $(TEST_IMAGE) go test -v ./internal/jobs/tiktok_test.go ./internal/jobs/jobs_suite_test.go
 
 test-reddit: docker-build-test
 	@docker run --user root $(ENV_FILE_ARG) -v $(PWD)/.masa:/home/masa -v $(PWD)/coverage:/app/coverage --rm --workdir /app -e DATA_DIR=/home/masa $(TEST_IMAGE) go test -v ./internal/jobs/reddit_test.go ./internal/jobs/redditapify/client_test.go ./api/types/reddit/reddit_suite_test.go

internal/capabilities/capabilities_suite_test.go

@@ -0,0 +1,13 @@
+package capabilities_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestCapabilities(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Capabilities Suite")
+}

internal/capabilities/detector.go

@@ -4,10 +4,14 @@ import (
 	"slices"
 	"strings"
 
+	"maps"
+
+	util "github.com/masa-finance/tee-types/pkg/util"
 	teetypes "github.com/masa-finance/tee-types/types"
 	"github.com/masa-finance/tee-worker/api/types"
 	"github.com/masa-finance/tee-worker/internal/jobs/twitter"
-	"maps"
+	"github.com/masa-finance/tee-worker/pkg/client"
+	"github.com/sirupsen/logrus"
 )
 
 // JobServerInterface defines the methods we need from JobServer to avoid circular dependencies
@@ -37,7 +41,7 @@ func DetectCapabilities(jc types.JobConfiguration, jobServer JobServerInterface)
 
 	hasAccounts := len(accounts) > 0
 	hasApiKeys := len(apiKeys) > 0
-	hasApifyKey := apifyApiKey != ""
+	hasApifyKey := hasValidApifyKey(apifyApiKey)
 
 	// Add Twitter-specific capabilities based on available authentication
 	if hasAccounts {
@@ -62,9 +66,17 @@ func DetectCapabilities(jc types.JobConfiguration, jobServer JobServerInterface)
 	if hasApifyKey {
 		capabilities[teetypes.TwitterApifyJob] = teetypes.TwitterApifyCaps
 		capabilities[teetypes.RedditJob] = teetypes.RedditCaps
+
+		// Merge TikTok search caps with any existing
+		existing := capabilities[teetypes.TiktokJob]
+		s := util.NewSet(existing...)
+		s.Add(teetypes.TiktokSearchCaps...)
+		capabilities[teetypes.TiktokJob] = s.Items()
+
 	}
 
 	// Add general TwitterJob capability if any Twitter auth is available
+	// TODO: this will get cleaned up with unique twitter capabilities
 	if hasAccounts || hasApiKeys || hasApifyKey {
 		var twitterJobCaps []teetypes.Capability
 		// Use the most comprehensive capabilities available
@@ -123,3 +135,25 @@ func parseApiKeys(apiKeys []string) []*twitter.TwitterApiKey {
 	}
 	return result
 }
+
+// hasValidApifyKey checks if the provided Apify API key is valid by attempting to validate it
+func hasValidApifyKey(apifyApiKey string) bool {
+	if apifyApiKey == "" {
+		return false
+	}
+
+	// Create temporary Apify client and validate the key
+	apifyClient, err := client.NewApifyClient(apifyApiKey)
+	if err != nil {
+		logrus.Errorf("Failed to create Apify client during capability detection: %v", err)
+		return false
+	}
+
+	if err := apifyClient.ValidateApiKey(); err != nil {
+		logrus.Errorf("Apify API key validation failed during capability detection: %v", err)
+		return false
+	}
+
+	logrus.Infof("Apify API key validated successfully during capability detection")
+	return true
+}

internal/capabilities/detector_test.go

@@ -1,12 +1,15 @@
-package capabilities
+package capabilities_test
 
 import (
-	"reflect"
+	"os"
 	"slices"
-	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
 
 	teetypes "github.com/masa-finance/tee-types/types"
 	"github.com/masa-finance/tee-worker/api/types"
+	. "github.com/masa-finance/tee-worker/internal/capabilities"
 )
 
 // MockJobServer implements JobServerInterface for testing
@@ -18,159 +21,169 @@ func (m *MockJobServer) GetWorkerCapabilities() teetypes.WorkerCapabilities {
 	return m.capabilities
 }
 
-func TestDetectCapabilities(t *testing.T) {
-	tests := []struct {
-		name      string
-		jc        types.JobConfiguration
-		jobServer JobServerInterface
-		expected  teetypes.WorkerCapabilities
-	}{
-		{
-			name: "With JobServer - gets capabilities from workers",
-			jc:   types.JobConfiguration{},
-			jobServer: &MockJobServer{
+var _ = Describe("DetectCapabilities", func() {
+	DescribeTable("capability detection scenarios",
+		func(jc types.JobConfiguration, jobServer JobServerInterface, expected teetypes.WorkerCapabilities) {
+			got := DetectCapabilities(jc, jobServer)
+
+			// Extract job type keys and sort for consistent comparison
+			gotKeys := make([]string, 0, len(got))
+			for jobType := range got {
+				gotKeys = append(gotKeys, jobType.String())
+			}
+
+			expectedKeys := make([]string, 0, len(expected))
+			for jobType := range expected {
+				expectedKeys = append(expectedKeys, jobType.String())
+			}
+
+			// Sort both slices for comparison
+			slices.Sort(gotKeys)
+			slices.Sort(expectedKeys)
+
+			// Compare the sorted slices
+			Expect(gotKeys).To(Equal(expectedKeys))
+		},
+		Entry("With JobServer - gets capabilities from workers",
+			types.JobConfiguration{},
+			&MockJobServer{
 				capabilities: teetypes.WorkerCapabilities{
 					teetypes.WebJob:       {teetypes.CapScraper},
 					teetypes.TelemetryJob: {teetypes.CapTelemetry},
 					teetypes.TiktokJob:    {teetypes.CapTranscription},
 					teetypes.TwitterJob:   {teetypes.CapSearchByQuery, teetypes.CapGetById, teetypes.CapGetProfileById},
 				},
 			},
-			expected: teetypes.WorkerCapabilities{
+			teetypes.WorkerCapabilities{
 				teetypes.WebJob:       {teetypes.CapScraper},
 				teetypes.TelemetryJob: {teetypes.CapTelemetry},
 				teetypes.TiktokJob:    {teetypes.CapTranscription},
 				teetypes.TwitterJob:   {teetypes.CapSearchByQuery, teetypes.CapGetById, teetypes.CapGetProfileById},
 			},
-		},
-		{
-			name:      "Without JobServer - basic capabilities only",
-			jc:        types.JobConfiguration{},
-			jobServer: nil,
-			expected: teetypes.WorkerCapabilities{
+		),
+		Entry("Without JobServer - basic capabilities only",
+			types.JobConfiguration{},
+			nil,
+			teetypes.WorkerCapabilities{
 				teetypes.WebJob:       {teetypes.CapScraper},
 				teetypes.TelemetryJob: {teetypes.CapTelemetry},
 				teetypes.TiktokJob:    {teetypes.CapTranscription},
 			},
-		},
-		{
-			name: "With Twitter accounts - adds credential capabilities",
-			jc: types.JobConfiguration{
+		),
+		Entry("With Twitter accounts - adds credential capabilities",
+			types.JobConfiguration{
 				"twitter_accounts": []string{"account1", "account2"},
 			},
-			jobServer: nil,
-			expected: teetypes.WorkerCapabilities{
+			nil,
+			teetypes.WorkerCapabilities{
 				teetypes.WebJob:               {teetypes.CapScraper},
 				teetypes.TelemetryJob:         {teetypes.CapTelemetry},
 				teetypes.TiktokJob:            {teetypes.CapTranscription},
 				teetypes.TwitterCredentialJob: teetypes.TwitterCredentialCaps,
 				teetypes.TwitterJob:           teetypes.TwitterCredentialCaps,
 			},
-		},
-		{
-			name: "With Twitter API keys - adds API capabilities",
-			jc: types.JobConfiguration{
+		),
+		Entry("With Twitter API keys - adds API capabilities",
+			types.JobConfiguration{
 				"twitter_api_keys": []string{"key1", "key2"},
 			},
-			jobServer: nil,
-			expected: teetypes.WorkerCapabilities{
+			nil,
+			teetypes.WorkerCapabilities{
 				teetypes.WebJob:        {teetypes.CapScraper},
 				teetypes.TelemetryJob:  {teetypes.CapTelemetry},
 				teetypes.TiktokJob:     {teetypes.CapTranscription},
 				teetypes.TwitterApiJob: teetypes.TwitterAPICaps,
 				teetypes.TwitterJob:    teetypes.TwitterAPICaps,
 			},
-		},
-		{
-			name: "With mock elevated Twitter API keys - only basic capabilities detected",
-			jc: types.JobConfiguration{
+		),
+		Entry("With mock elevated Twitter API keys - only basic capabilities detected",
+			types.JobConfiguration{
 				"twitter_api_keys": []string{"Bearer abcd1234-ELEVATED"},
 			},
-			jobServer: nil,
-			expected: teetypes.WorkerCapabilities{
+			nil,
+			teetypes.WorkerCapabilities{
 				teetypes.WebJob:       {teetypes.CapScraper},
 				teetypes.TelemetryJob: {teetypes.CapTelemetry},
 				teetypes.TiktokJob:    {teetypes.CapTranscription},
 				// Note: Mock elevated keys will be detected as basic since we can't make real API calls in tests
 				teetypes.TwitterApiJob: teetypes.TwitterAPICaps,
 				teetypes.TwitterJob:    teetypes.TwitterAPICaps,
 			},
-		},
-	}
+		),
+	)
 
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got := DetectCapabilities(tt.jc, tt.jobServer)
+	Context("Scraper Types", func() {
+		DescribeTable("scraper type detection",
+			func(jc types.JobConfiguration, expectedKeys []string) {
+				caps := DetectCapabilities(jc, nil)
 
-			// Extract job type keys and sort for consistent comparison
-			gotKeys := make([]string, 0, len(got))
-			for jobType := range got {
-				gotKeys = append(gotKeys, jobType.String())
-			}
+				jobNames := make([]string, 0, len(caps))
+				for jobType := range caps {
+					jobNames = append(jobNames, jobType.String())
+				}
 
-			expectedKeys := make([]string, 0, len(tt.expected))
-			for jobType := range tt.expected {
-				expectedKeys = append(expectedKeys, jobType.String())
-			}
+				// Sort both slices for comparison
+				slices.Sort(jobNames)
+				expectedSorted := make([]string, len(expectedKeys))
+				copy(expectedSorted, expectedKeys)
+				slices.Sort(expectedSorted)
 
-			// Sort both slices for comparison
-			slices.Sort(gotKeys)
-			slices.Sort(expectedKeys)
+				// Compare the sorted slices
+				Expect(jobNames).To(Equal(expectedSorted))
+			},
+			Entry("Basic scrapers only",
+				types.JobConfiguration{},
+				[]string{"web", "telemetry", "tiktok"},
+			),
+			Entry("With Twitter accounts",
+				types.JobConfiguration{
+					"twitter_accounts": []string{"user1:pass1"},
+				},
+				[]string{"web", "telemetry", "tiktok", "twitter", "twitter-credential"},
+			),
+			Entry("With Twitter API keys",
+				types.JobConfiguration{
+					"twitter_api_keys": []string{"key1"},
+				},
+				[]string{"web", "telemetry", "tiktok", "twitter", "twitter-api"},
+			),
+		)
+	})
 
-			// Compare the sorted slices
-			if !reflect.DeepEqual(gotKeys, expectedKeys) {
-				t.Errorf("DetectCapabilities() job types = %v, want %v", gotKeys, expectedKeys)
+	Context("Apify Integration", func() {
+		It("should add enhanced capabilities when valid Apify API key is provided", func() {
+			apifyKey := os.Getenv("APIFY_API_KEY")
+			if apifyKey == "" {
+				Skip("APIFY_API_KEY is not set")
 			}
-		})
-	}
-}
 
-func TestDetectCapabilities_ScraperTypes(t *testing.T) {
-	tests := []struct {
-		name         string
-		jc           types.JobConfiguration
-		expectedKeys []string // scraper names we expect
-	}{
-		{
-			name:         "Basic scrapers only",
-			jc:           types.JobConfiguration{},
-			expectedKeys: []string{"web", "telemetry", "tiktok"},
-		},
-		{
-			name: "With Twitter accounts",
-			jc: types.JobConfiguration{
-				"twitter_accounts": []string{"user1:pass1"},
-			},
-			expectedKeys: []string{"web", "telemetry", "tiktok", "twitter", "twitter-credential"},
-		},
-		{
-			name: "With Twitter API keys",
-			jc: types.JobConfiguration{
-				"twitter_api_keys": []string{"key1"},
-			},
-			expectedKeys: []string{"web", "telemetry", "tiktok", "twitter", "twitter-api"},
-		},
-	}
+			jc := types.JobConfiguration{
+				"apify_api_key": apifyKey,
+			}
 
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			caps := DetectCapabilities(tt.jc, nil)
+			caps := DetectCapabilities(jc, nil)
 
-			jobNames := make([]string, 0, len(caps))
-			for jobType := range caps {
-				jobNames = append(jobNames, jobType.String())
-			}
+			// TikTok should gain search capabilities with valid key
+			tiktokCaps, ok := caps[teetypes.TiktokJob]
+			Expect(ok).To(BeTrue(), "expected tiktok capabilities to be present")
+			Expect(tiktokCaps).To(ContainElement(teetypes.CapSearchByQuery), "expected tiktok to include CapSearchByQuery capability")
+			Expect(tiktokCaps).To(ContainElement(teetypes.CapSearchByTrending), "expected tiktok to include CapSearchByTrending capability")
 
-			// Sort both slices for comparison
-			slices.Sort(jobNames)
-			expectedSorted := make([]string, len(tt.expectedKeys))
-			copy(expectedSorted, tt.expectedKeys)
-			slices.Sort(expectedSorted)
+			// Twitter-Apify job should be present with follower/following capabilities
+			twitterApifyCaps, ok := caps[teetypes.TwitterApifyJob]
+			Expect(ok).To(BeTrue(), "expected twitter-apify capabilities to be present")
+			Expect(twitterApifyCaps).To(ContainElement(teetypes.CapGetFollowers), "expected twitter-apify to include CapGetFollowers capability")
+			Expect(twitterApifyCaps).To(ContainElement(teetypes.CapGetFollowing), "expected twitter-apify to include CapGetFollowing capability")
 
-			// Compare the sorted slices
-			if !reflect.DeepEqual(jobNames, expectedSorted) {
-				t.Errorf("Expected capabilities %v, got %v", expectedSorted, jobNames)
-			}
+			// Reddit should be present
+			_, hasReddit := caps[teetypes.RedditJob]
+			Expect(hasReddit).To(BeTrue(), "expected reddit capabilities to be present")
 		})
-	}
+	})
+})
+
+// Helper function to check if a job type exists in capabilities
+func hasJobType(capabilities teetypes.WorkerCapabilities, jobName string) bool {
+	_, exists := capabilities[teetypes.JobType(jobName)]
+	return exists
 }