jami-daemon/test/unitTest/media/test_media_encoder.cpp

/*
 *  Copyright (C) 2018-2019 Savoir-faire Linux Inc.
 *
 *  Author: Philippe Gorley <philippe.gorley@savoirfairelinux.com>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA.
 */

#include <cppunit/TestAssert.h>
#include <cppunit/TestFixture.h>
#include <cppunit/extensions/HelperMacros.h>

#include "dring.h"
#include "fileutils.h"
#include "libav_deps.h"
#include "media_encoder.h"
#include "media_io_handle.h"
#include "system_codec_container.h"

#include "../../test_runner.h"

namespace ring { namespace test {

class MediaEncoderTest : public CppUnit::TestFixture {
public:
    static std::string name() { return "media_encoder"; }

    void setUp();
    void tearDown();

private:
    void testMultiStream();

    CPPUNIT_TEST_SUITE(MediaEncoderTest);
    CPPUNIT_TEST(testMultiStream);
    CPPUNIT_TEST_SUITE_END();

    std::unique_ptr<MediaEncoder> encoder_;
    std::vector<std::string> files_;
};

CPPUNIT_TEST_SUITE_NAMED_REGISTRATION(MediaEncoderTest, MediaEncoderTest::name());

void
MediaEncoderTest::setUp()
{
    DRing::init(DRing::InitFlag(DRing::DRING_FLAG_DEBUG | DRing::DRING_FLAG_CONSOLE_LOG));
    libav_utils::ring_avcodec_init();
    encoder_.reset(new MediaEncoder);
    files_.push_back("test.mkv");
}

void
MediaEncoderTest::tearDown()
{
    // clean up behind ourselves
    for (const auto& file : files_)
        fileutils::remove(file);
    DRing::fini();
}

static AVFrame*
getVideoFrame(int width, int height, int frame_index)
{
    int x, y;
    AVFrame* frame = av_frame_alloc();
    if (!frame)
        return nullptr;

    frame->format = AV_PIX_FMT_YUV420P;
    frame->width = width;
    frame->height = height;

    if (av_frame_get_buffer(frame, 32) < 0) {
        av_frame_free(&frame);
        return nullptr;
    }

    /* Y */
    for (y = 0; y < height; y++)
        for (x = 0; x < width; x++)
            frame->data[0][y * frame->linesize[0] + x] = x + y + frame_index * 3;

    /* Cb and Cr */
    for (y = 0; y < height / 2; y++) {
        for (x = 0; x < width / 2; x++) {
            frame->data[1][y * frame->linesize[1] + x] = 128 + y + frame_index * 2;
            frame->data[2][y * frame->linesize[2] + x] = 64 + x + frame_index * 5;
        }
    }

    return frame;
}

static AVFrame*
getAudioFrame(int sampleRate, int nbSamples, int nbChannels)
{
    const constexpr float pi = 3.14159265358979323846264338327950288; // M_PI
    const float tincr = 2 * pi * 440.0 / sampleRate;
    float t = 0;
    AVFrame* frame = av_frame_alloc();
    if (!frame)
        return nullptr;

    frame->format = AV_SAMPLE_FMT_S16;
    frame->channels = nbChannels;
    frame->channel_layout = av_get_default_channel_layout(nbChannels);
    frame->nb_samples = nbSamples;
    frame->sample_rate = sampleRate;

    if (av_frame_get_buffer(frame, 0) < 0) {
        av_frame_free(&frame);
        return nullptr;
    }

    auto samples = reinterpret_cast<uint16_t*>(frame->data[0]);
    for (int i = 0; i < 200; ++i) {
        for (int j = 0; j < nbSamples; ++j) {
            samples[2 * j] = static_cast<int>(sin(t) * 10000);
            for (int k = 1; k < nbChannels; ++k) {
                samples[2 * j + k] = samples[2 * j];
            }
            t += tincr;
        }
    }

    return frame;
}

void
MediaEncoderTest::testMultiStream()
{
    const constexpr int sampleRate = 48000;
    const constexpr int nbChannels = 2;
    const constexpr int width = 320;
    const constexpr int height = 240;
    std::map<std::string, std::string> options;
    options["title"] = "Encoder Unit Test";
    options["description"] = "Description goes here";
    options["sample_rate"] = std::to_string(sampleRate);
    options["channels"] = std::to_string(nbChannels);
    options["width"] = std::to_string(width);
    options["height"] = std::to_string(height);
    auto vp8Codec = std::static_pointer_cast<ring::SystemVideoCodecInfo>(
        getSystemCodecContainer()->searchCodecByName("VP8", ring::MEDIA_VIDEO)
    );
    auto opusCodec = std::static_pointer_cast<SystemAudioCodecInfo>(
        getSystemCodecContainer()->searchCodecByName("opus", ring::MEDIA_AUDIO)
    );

    try {
        encoder_->openOutput("test.mkv");
        encoder_->setOptions(options);
        int videoIdx = encoder_->addStream(*vp8Codec.get());
        CPPUNIT_ASSERT(videoIdx >= 0);
        CPPUNIT_ASSERT(encoder_->getStreamCount() == 1);
        int audioIdx = encoder_->addStream(*opusCodec.get());
        CPPUNIT_ASSERT(audioIdx >= 0);
        CPPUNIT_ASSERT(videoIdx != audioIdx);
        CPPUNIT_ASSERT(encoder_->getStreamCount() == 2);
        encoder_->setIOContext(nullptr);
        encoder_->startIO();
        int sentSamples = 0;
        AVFrame* audio = nullptr;
        AVFrame* video = nullptr;
        for (int i = 0; i < 25; ++i) {
            audio = getAudioFrame(sampleRate, 0.02*sampleRate, nbChannels);
            CPPUNIT_ASSERT(audio);
            audio->pts = sentSamples;
            video = getVideoFrame(width, height, i);
            CPPUNIT_ASSERT(video);
            video->pts = i;

            CPPUNIT_ASSERT(encoder_->encode(audio, audioIdx) >= 0);
            sentSamples += audio->nb_samples;
            CPPUNIT_ASSERT(encoder_->encode(video, videoIdx) >= 0);

            av_frame_free(&audio);
            av_frame_free(&video);
        }
        CPPUNIT_ASSERT(encoder_->flush() >= 0);
    } catch (const MediaEncoderException& e) {
        CPPUNIT_FAIL(e.what());
    }
}

}} // namespace ring::test

RING_TEST_RUNNER(ring::test::MediaEncoderTest::name());