Skip to content

Commit 107ccc4

Browse files
committed
Add tests for custom where expression in Postgres Store
This adds comprehensive test coverage for the new custom where expression feature merged in PR #235: - Test custom where expression without maxScore - Test custom where expression combined with maxScore - Test custom where expression with custom parameters
1 parent 3394f86 commit 107ccc4

File tree

1 file changed

+131
-0
lines changed

1 file changed

+131
-0
lines changed

src/store/tests/Bridge/Postgres/StoreTest.php

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -461,4 +461,135 @@ public function testQueryWithNullMetadata()
461461
$this->assertCount(1, $results);
462462
$this->assertSame([], $results[0]->metadata->getArrayCopy());
463463
}
464+
465+
public function testQueryWithCustomWhereExpression()
466+
{
467+
$pdo = $this->createMock(\PDO::class);
468+
$statement = $this->createMock(\PDOStatement::class);
469+
470+
$store = new Store($pdo, 'embeddings_table', 'embedding');
471+
472+
$expectedSql = 'SELECT id, embedding AS embedding, metadata, (embedding <-> :embedding) AS score
473+
FROM embeddings_table
474+
WHERE metadata->>\'category\' = \'products\'
475+
ORDER BY score ASC
476+
LIMIT 5';
477+
478+
$pdo->expects($this->once())
479+
->method('prepare')
480+
->with($this->callback(function ($sql) use ($expectedSql) {
481+
return $this->normalizeQuery($sql) === $this->normalizeQuery($expectedSql);
482+
}))
483+
->willReturn($statement);
484+
485+
$statement->expects($this->once())
486+
->method('execute')
487+
->with(['embedding' => '[0.1,0.2,0.3]']);
488+
489+
$statement->expects($this->once())
490+
->method('fetchAll')
491+
->with(\PDO::FETCH_ASSOC)
492+
->willReturn([]);
493+
494+
$results = $store->query(new Vector([0.1, 0.2, 0.3]), ['where' => 'metadata->>\'category\' = \'products\'']);
495+
496+
$this->assertCount(0, $results);
497+
}
498+
499+
public function testQueryWithCustomWhereExpressionAndMaxScore()
500+
{
501+
$pdo = $this->createMock(\PDO::class);
502+
$statement = $this->createMock(\PDOStatement::class);
503+
504+
$store = new Store($pdo, 'embeddings_table', 'embedding');
505+
506+
$expectedSql = 'SELECT id, embedding AS embedding, metadata, (embedding <-> :embedding) AS score
507+
FROM embeddings_table
508+
WHERE (embedding <-> :embedding) <= :maxScore AND (metadata->>\'active\' = \'true\')
509+
ORDER BY score ASC
510+
LIMIT 5';
511+
512+
$pdo->expects($this->once())
513+
->method('prepare')
514+
->with($this->callback(function ($sql) use ($expectedSql) {
515+
return $this->normalizeQuery($sql) === $this->normalizeQuery($expectedSql);
516+
}))
517+
->willReturn($statement);
518+
519+
$statement->expects($this->once())
520+
->method('execute')
521+
->with([
522+
'embedding' => '[0.1,0.2,0.3]',
523+
'maxScore' => 0.5,
524+
]);
525+
526+
$statement->expects($this->once())
527+
->method('fetchAll')
528+
->with(\PDO::FETCH_ASSOC)
529+
->willReturn([]);
530+
531+
$results = $store->query(new Vector([0.1, 0.2, 0.3]), [
532+
'maxScore' => 0.5,
533+
'where' => 'metadata->>\'active\' = \'true\'',
534+
]);
535+
536+
$this->assertCount(0, $results);
537+
}
538+
539+
public function testQueryWithCustomWhereExpressionAndParams()
540+
{
541+
$pdo = $this->createMock(\PDO::class);
542+
$statement = $this->createMock(\PDOStatement::class);
543+
544+
$store = new Store($pdo, 'embeddings_table', 'embedding');
545+
546+
$expectedSql = 'SELECT id, embedding AS embedding, metadata, (embedding <-> :embedding) AS score
547+
FROM embeddings_table
548+
WHERE metadata->>\'crawlId\' = :crawlId AND id != :currentId
549+
ORDER BY score ASC
550+
LIMIT 5';
551+
552+
$pdo->expects($this->once())
553+
->method('prepare')
554+
->with($this->callback(function ($sql) use ($expectedSql) {
555+
return $this->normalizeQuery($sql) === $this->normalizeQuery($expectedSql);
556+
}))
557+
->willReturn($statement);
558+
559+
$uuid = Uuid::v4();
560+
$crawlId = '396af6fe-0dfd-47ed-b222-3dbcced3f38e';
561+
562+
$statement->expects($this->once())
563+
->method('execute')
564+
->with([
565+
'embedding' => '[0.1,0.2,0.3]',
566+
'crawlId' => $crawlId,
567+
'currentId' => $uuid->toRfc4122(),
568+
]);
569+
570+
$statement->expects($this->once())
571+
->method('fetchAll')
572+
->with(\PDO::FETCH_ASSOC)
573+
->willReturn([
574+
[
575+
'id' => Uuid::v4()->toRfc4122(),
576+
'embedding' => '[0.4,0.5,0.6]',
577+
'metadata' => json_encode(['crawlId' => $crawlId, 'url' => 'https://example.com']),
578+
'score' => 0.85,
579+
],
580+
]);
581+
582+
$results = $store->query(new Vector([0.1, 0.2, 0.3]), [
583+
'where' => 'metadata->>\'crawlId\' = :crawlId AND id != :currentId',
584+
'params' => [
585+
'crawlId' => $crawlId,
586+
'currentId' => $uuid->toRfc4122(),
587+
],
588+
]);
589+
590+
$this->assertCount(1, $results);
591+
$this->assertSame(0.85, $results[0]->score);
592+
$this->assertSame($crawlId, $results[0]->metadata['crawlId']);
593+
$this->assertSame('https://example.com', $results[0]->metadata['url']);
594+
}
464595
}

0 commit comments

Comments
 (0)