Examples with QueryAssertions io.trino.sql.query.QueryAssertions used on opensource projects

Example 1 with QueryAssertions

use of io.trino.sql.query.QueryAssertions in project trino by trinodb.

the class TestOperators method init.

@BeforeClass
public void init() {
    Session session = testSessionBuilder().setTimeZoneKey(TestingSession.DEFAULT_TIME_ZONE_KEY).build();
    assertions = new QueryAssertions(session);
}

Example 2 with QueryAssertions

use of io.trino.sql.query.QueryAssertions in project trino by trinodb.

the class TestKafkaWithConfluentSchemaRegistryMinimalFunctionality method assertTopic.

private void assertTopic(TestingKafka testingKafka, String topicName, String initialQuery, String evolvedQuery, boolean isKeyIncluded, Map<String, String> producerConfig) {
    assertNotExists(topicName);
    List<ProducerRecord<Long, GenericRecord>> messages = createMessages(topicName, MESSAGE_COUNT, true);
    testingKafka.sendMessages(messages.stream(), producerConfig);
    waitUntilTableExists(topicName);
    assertCount(topicName, MESSAGE_COUNT);
    QueryAssertions queryAssertions = new QueryAssertions(getQueryRunner());
    queryAssertions.query(initialQuery).assertThat().containsAll(getExpectedValues(messages, INITIAL_SCHEMA, isKeyIncluded));
    List<ProducerRecord<Long, GenericRecord>> newMessages = createMessages(topicName, MESSAGE_COUNT, false);
    testingKafka.sendMessages(newMessages.stream(), producerConfig);
    List<ProducerRecord<Long, GenericRecord>> allMessages = ImmutableList.<ProducerRecord<Long, GenericRecord>>builder().addAll(messages).addAll(newMessages).build();
    assertCount(topicName, allMessages.size());
    queryAssertions.query(evolvedQuery).assertThat().containsAll(getExpectedValues(messages, EVOLVED_SCHEMA, isKeyIncluded));
}

Example 3 with QueryAssertions

use of io.trino.sql.query.QueryAssertions in project trino by trinodb.

the class TestRevokeOnTable method initClass.

@BeforeClass
public void initClass() throws Exception {
    SchemaTableName table = new SchemaTableName("default", "table_one");
    queryRunner = DistributedQueryRunner.builder(userWithAllPrivileges).build();
    Grants<SchemaTableName> tableGrants = new MutableGrants<>();
    tableGrants.grant(new TrinoPrincipal(USER, admin.getUser()), table, EnumSet.allOf(Privilege.class), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithAllPrivileges.getUser()), table, EnumSet.allOf(Privilege.class), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithCreate.getUser()), table, ImmutableSet.of(Privilege.CREATE), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithSelect.getUser()), table, ImmutableSet.of(Privilege.SELECT), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithInsert.getUser()), table, ImmutableSet.of(Privilege.INSERT), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithUpdate.getUser()), table, ImmutableSet.of(Privilege.UPDATE), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithDelete.getUser()), table, ImmutableSet.of(Privilege.DELETE), true);
    MockConnectorFactory connectorFactory = MockConnectorFactory.builder().withListSchemaNames(session -> ImmutableList.of("default")).withListTables((session, schemaName) -> "default".equalsIgnoreCase(schemaName) ? ImmutableList.of(table) : ImmutableList.of()).withGetTableHandle((session, tableName) -> tableName.equals(table) ? new MockConnectorTableHandle(tableName) : null).withSchemaGrants(new MutableGrants<>()).withTableGrants(tableGrants).build();
    queryRunner.installPlugin(new MockConnectorPlugin(connectorFactory));
    queryRunner.createCatalog("local", "mock");
    assertions = new QueryAssertions(queryRunner);
}

Example 4 with QueryAssertions

use of io.trino.sql.query.QueryAssertions in project trino by trinodb.

the class SqlDataTypeTest method verifyPredicate.

private void verifyPredicate(QueryRunner queryRunner, Session session, TestTable testTable) {
    String queryWithAll = "SELECT 'all found' FROM " + testTable.getName() + " WHERE " + IntStream.range(0, testCases.size()).mapToObj(this::getPredicate).collect(joining(" AND "));
    MaterializedResult result = queryRunner.execute(session, queryWithAll);
    if (result.getOnlyColumnAsSet().equals(Set.of("all found"))) {
        return;
    }
    // Closing QueryAssertions would close the QueryRunner
    @SuppressWarnings("resource") QueryAssertions queryAssertions = new QueryAssertions(queryRunner);
    for (int column = 0; column < testCases.size(); column++) {
        assertThat(queryAssertions.query(session, "SELECT 'found' FROM " + testTable.getName() + " WHERE " + getPredicate(column))).matches("VALUES 'found'");
    }
}

Example 5 with QueryAssertions

use of io.trino.sql.query.QueryAssertions in project trino by trinodb.

the class SqlDataTypeTest method verifySelect.

private void verifySelect(QueryRunner queryRunner, Session session, TestTable testTable) {
    // Closing QueryAssertions would close the QueryRunner
    @SuppressWarnings("resource") QueryAssertions queryAssertions = new QueryAssertions(queryRunner);
    QueryAssert assertion = assertThat(queryAssertions.query(session, "SELECT * FROM " + testTable.getName()));
    MaterializedResult expected = queryRunner.execute(session, testCases.stream().map(TestCase::getExpectedLiteral).collect(joining(",", "VALUES ROW(", ")")));
    // Verify types if specified
    for (int column = 0; column < testCases.size(); column++) {
        TestCase testCase = testCases.get(column);
        if (testCase.getExpectedType().isPresent()) {
            Type expectedType = testCase.getExpectedType().get();
            assertion.outputHasType(column, expectedType);
            assertThat(expected.getTypes()).as(format("Expected literal type at column %d (check consistency of expected type and expected literal)", column + 1)).element(column).isEqualTo(expectedType);
        }
    }
    assertion.matches(expected);
}