Fix test

DataDog · Jan 22, 2025 · eaed1bb · eaed1bb
1 parent 6a5371d
commit eaed1bb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/test/fakeintake/server/serverstore/db.go b/test/fakeintake/server/serverstore/db.go
@@ -100,6 +100,7 @@ func newSQLStore() *sqlStore {
 	CREATE TABLE IF NOT EXISTS payloads (
 		id INTEGER PRIMARY KEY AUTOINCREMENT,
 		timestamp INTEGER NOT NULL,
+		apiKey VARCHAR(32) NOT NULL,
 		data BLOB NOT NULL,
 		encoding VARCHAR(10) NOT NULL,
 		content_type VARCHAR(20),
@@ -129,7 +130,7 @@ func (s *sqlStore) MostRecentPayloadAPIKey(route string) (string, error) {
 // AppendPayload adds a payload to the store and tries parsing and adding a dumped json to the parsed store
 func (s *sqlStore) AppendPayload(route string, apiKey string, data []byte, encoding string, contentType string, collectTime time.Time) error {
 	now := time.Now()
-	_, err := s.db.Exec("INSERT INTO payloads (timestamp, apiKey, data, encoding, content_type, route) VALUES (?, ?, ?, ?, ?, ?)", collectTime.Unix(), data, encoding, contentType, route)
+	_, err := s.db.Exec("INSERT INTO payloads (timestamp, apiKey, data, encoding, content_type, route) VALUES (?, ?, ?, ?, ?, ?)", collectTime.Unix(), apiKey, data, encoding, contentType, route)
 	if err != nil {
 		return err
 	}
@@ -159,7 +160,7 @@ func (s *sqlStore) CleanUpPayloadsOlderThan(time time.Time) {
 // GetRawPayloads returns all raw payloads for a given route
 func (s *sqlStore) GetRawPayloads(route string) []api.Payload {
 	now := time.Now()
-	rows, err := s.db.Query("SELECT timestamp, data, encoding, content_type FROM payloads WHERE route = ?", route)
+	rows, err := s.db.Query("SELECT timestamp, apiKey, data, encoding, content_type FROM payloads WHERE route = ?", route)
 	if err != nil {
 		log.Println("Error fetching raw payloads: ", err)
 		return nil
@@ -168,18 +169,20 @@ func (s *sqlStore) GetRawPayloads(route string) []api.Payload {
 	s.metrics.readLatency.WithLabelValues(route).Observe(time.Since(now).Seconds())
 
 	var timestamp int64
+	var apiKey string
 	var data []byte
 	var encoding string
 	var contentType string
 	payloads := []api.Payload{}
 	for rows.Next() {
-		err := rows.Scan(&timestamp, &data, &encoding, &contentType)
+		err := rows.Scan(&timestamp, &apiKey, &data, &encoding, &contentType)
 		if err != nil {
 			log.Println("Error scanning raw payload: ", err)
 			continue
 		}
 		payloads = append(payloads, api.Payload{
 			Timestamp:   time.Unix(timestamp, 0),
+			APIKey:      apiKey,
 			Data:        data,
 			Encoding:    encoding,
 			ContentType: contentType,

diff --git a/test/fakeintake/server/serverstore/store_test.go b/test/fakeintake/server/serverstore/store_test.go
@@ -37,12 +37,13 @@ func (suite *StoreTestSuite) TestAppendPayload() {
 	data := []byte(`{"key":"value"}`)
 	parserMap["testRoute"] = jsonParser
 	defer delete(parserMap, "testRoute")
-	err := store.AppendPayload("testRoute", data, "json", "", time.Now())
+	err := store.AppendPayload("testRoute", "1234", data, "json", "", time.Now())
 	assert.NoError(suite.T(), err)
 
 	rawPayloads := store.GetRawPayloads("testRoute")
 	assert.Len(suite.T(), rawPayloads, 1)
 	assert.Equal(suite.T(), data, rawPayloads[0].Data)
+	assert.Equal(suite.T(), "1234", rawPayloads[0].APIKey)
 
 	jsonPayloads, err := GetJSONPayloads(store, "testRoute")
 	require.NoError(suite.T(), err)
@@ -59,10 +60,10 @@ func (suite *StoreTestSuite) TestCleanUpPayloadsOlderThan() {
 	parserMap["testRoute"] = jsonParser
 	defer delete(parserMap, "testRoute")
 	// Add an old payload expected to be cleaned up first
-	err := store.AppendPayload("testRoute", []byte("{}"), "json", "", now.Add(-48*time.Hour))
+	err := store.AppendPayload("testRoute", "1234", []byte("{}"), "json", "", now.Add(-48*time.Hour))
 	require.NoError(suite.T(), err)
 
-	err = store.AppendPayload("testRoute", []byte("{}"), "json", "", now)
+	err = store.AppendPayload("testRoute", "1234", []byte("{}"), "json", "", now)
 	require.NoError(suite.T(), err)
 
 	rawPayloads := store.GetRawPayloads("testRoute")
@@ -82,10 +83,10 @@ func (suite *StoreTestSuite) TestGetRouteStats() {
 	store := suite.StoreConstructor()
 	defer store.Close()
 
-	err := store.AppendPayload("routeA", []byte("{}"), "json", "", time.Now())
+	err := store.AppendPayload("routeA", "1234", []byte("{}"), "json", "", time.Now())
 	require.NoError(suite.T(), err)
 
-	err = store.AppendPayload("routeB", []byte("{}"), "json", "", time.Now())
+	err = store.AppendPayload("routeB", "1234", []byte("{}"), "json", "", time.Now())
 	require.NoError(suite.T(), err)
 
 	stats := store.GetRouteStats()
@@ -100,7 +101,7 @@ func (suite *StoreTestSuite) TestFlush() {
 
 	parserMap["testRoute"] = jsonParser
 	defer delete(parserMap, "testRoute")
-	err := store.AppendPayload("testRoute", []byte("{}"), "json", "", time.Now())
+	err := store.AppendPayload("testRoute", "1234", []byte("{}"), "json", "", time.Now())
 	require.NoError(suite.T(), err)
 
 	store.Flush()