<commit_before>package org.rakam.analysis;

import com.facebook.presto.sql.parser.SqlParser;
import com.facebook.presto.sql.tree.Expression;
import com.facebook.presto.sql.tree.Statement;
import org.testng.annotations.Test;

import java.util.stream.Collectors;

import static com.facebook.presto.sql.RakamExpressionFormatter.formatIdentifier;
import static com.facebook.presto.sql.RakamSqlFormatter.Formatter.formatQuery;
import static com.facebook.presto.sql.RakamSqlFormatter.formatExpression;
import static org.testng.Assert.assertEquals;

public class TestQueryFormatter
{
    @Test
    public void testSimpleExpression()
            throws Exception
    {
        Expression expression = new SqlParser().createExpression("test = 'test'");

        assertEquals(formatExpression(expression,
                name -> {
                    throw new UnsupportedOperationException();
                },
                name -> "\"dummy\".\"" + name + "\"", '"'), "(\"dummy\".\"test\" = 'test')");
    }

    @Test
    public void testSimpleQuery()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement("select * from testcollection");

        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"SELECT *\n" +
                "   FROM\n" +
                "     dummy");
    }

    @Test
    public void testJoinQuery()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement
                ("select * from testcollection join anothercollection on (anothercollection.test = testcollection.test)");

        // TODO: decide if we should also format expressions in QueryFormatter
        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"SELECT *\n" +
                        "   FROM\n" +
                        "     (dummy\n" +
                        "   INNER JOIN dummy ON (\"anothercollection\".\"test\" = \"testcollection\".\"test\"))");
    }

    @Test
    public void testQueryWithCTE()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement("with test as (select * from collection) select * from test");

        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"WITH\n" +
                "     \"test\" AS (\n" +
                "      SELECT *\n" +
                "      FROM\n" +
                "        dummy\n" +
                "   ) \n" +
                "   SELECT *\n" +
                "   FROM\n" +
                "     test");
    }

    @Test
    public void testQueryWithCTEDuplicateName()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement("with test as (select * from collection) select * from collection.test");

        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"WITH\n" +
                "    \"test\" AS (\n" +
                "      SELECT *\n" +
                "      FROM\n" +
                "        dummy\n" +
                "   ) \n" +
                "   SELECT *\n" +
                "   FROM\n" +
                "     dummy");
    }

    @Test
    public void testExpressionFormatterFormatTable()
            throws Exception
    {<commit_msg>Fix failing test<commit_after>package org.rakam.analysis;

import com.facebook.presto.sql.parser.SqlParser;
import com.facebook.presto.sql.tree.Expression;
import com.facebook.presto.sql.tree.Statement;
import org.testng.annotations.Test;

import java.util.stream.Collectors;

import static com.facebook.presto.sql.RakamExpressionFormatter.formatIdentifier;
import static com.facebook.presto.sql.RakamSqlFormatter.Formatter.formatQuery;
import static com.facebook.presto.sql.RakamSqlFormatter.formatExpression;
import static org.testng.Assert.assertEquals;

public class TestQueryFormatter
{
    @Test
    public void testSimpleExpression()
            throws Exception
    {
        Expression expression = new SqlParser().createExpression("test = 'test'");

        assertEquals(formatExpression(expression,
                name -> {
                    throw new UnsupportedOperationException();
                },
                name -> "\"dummy\".\"" + name + "\"", '"'), "(\"dummy\".\"test\" = 'test')");
    }

    @Test
    public void testSimpleQuery()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement("select * from testcollection");

        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"SELECT *\n" +
                "   FROM\n" +
                "     dummy");
    }

    @Test
    public void testJoinQuery()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement
                ("select * from testcollection join anothercollection on (anothercollection.test = testcollection.test)");

        // TODO: decide if we should also format expressions in QueryFormatter
        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"SELECT *\n" +
                        "   FROM\n" +
                        "     (dummy\n" +
                        "   INNER JOIN dummy ON (\"anothercollection\".\"test\" = \"testcollection\".\"test\"))");
    }

    @Test
    public void testQueryWithCTE()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement("with test as (select * from collection) select * from test");

        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"WITH\n" +
                "     \"test\" AS (\n" +
                "      SELECT *\n" +
                "      FROM\n" +
                "        dummy\n" +
                "   ) \n" +
                "   SELECT *\n" +
                "   FROM\n" +
                "     test");
    }

    @Test
    public void testQueryWithCTEDuplicateName()
            throws Exception
    {
        Statement statement = new SqlParser().createStatement("with test as (select * from collection) select * from collection.test");

        assertEquals(formatQuery(statement, name -> "dummy", '"').trim(),"WITH\n" +
                "     \"test\" AS (\n" +
                "      SELECT *\n" +
                "      FROM\n" +
                "        dummy\n" +
                "   ) \n" +
                "   SELECT *\n" +
                "   FROM\n" +
                "     dummy");
    }

    @Test
    public void testExpressionFormatterFormatTable()
            throws Exception
    {
{'fraction_non_alphanumeric': 0.1264687693259122, 'fraction_numerical': 0.0, 'mean_word_length': 1.7907679033649697, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6647619', 'n_tokens_mistral': 1727, 'n_tokens_neox': 1652, 'n_words': 366}